gemini-code-assist[bot] commented on code in PR #18774:
URL: https://github.com/apache/tvm/pull/18774#discussion_r2803259587
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -1797,7 +1821,9 @@ def _impl_v1(cls, bb, inputs, attr, params):
indices.append(index)
# When splits isnt specified divide evenly over axis.
else:
- indices = attr["tvm_custom"]["num_outputs"]
+ axis = attr.get("axis", 0)
+ num_outputs = attr["tvm_custom"]["num_outputs"]
+ indices = cls._compute_split_indices(inputs[0], axis, num_outputs)
return relax.op.split(inputs[0], indices, attr.get("axis", 0))
Review Comment:
To improve readability and avoid a redundant dictionary lookup, you can
reuse the `axis` variable that was already retrieved on line 1824.
```suggestion
return relax.op.split(inputs[0], indices, axis)
```
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -1818,7 +1844,9 @@ def _impl_v13(cls, bb, inputs, attr, params):
raise ValueError("Dynamic Split not yet supported")
# When splits isnt specified divide evenly over axis.
else:
- indices = attr["tvm_custom"]["num_outputs"]
+ axis = attr.get("axis", 0)
+ num_outputs = attr.get("num_outputs",
attr["tvm_custom"]["num_outputs"])
+ indices = cls._compute_split_indices(inputs[0], axis, num_outputs)
return relax.op.split(inputs[0], indices, attr.get("axis", 0))
Review Comment:
Similar to the change in `_impl_v1`, you can reuse the `axis` variable that
was retrieved on line 1847 to improve readability and avoid a redundant
dictionary lookup.
```suggestion
return relax.op.split(inputs[0], indices, axis)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
