gemini-code-assist[bot] commented on code in PR #18677:
URL: https://github.com/apache/tvm/pull/18677#discussion_r2713356700
##########
python/tvm/relax/transform/legalize_ops/nn.py:
##########
@@ -775,3 +775,8 @@ def nll_loss_without_weight(predictions, targets,
reduction, ignore_index):
reduction=call.attrs.reduction,
ignore_index=call.attrs.ignore_index,
)
+
+
+@register_legalize("relax.nn.batch_flatten")
+def _nn_batch_flatten(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(topi.reshape, call.args[0],
call.struct_info.shape.values)
Review Comment:
The legalization for `relax.nn.batch_flatten` assumes that
`call.struct_info.shape` is always available. However, if the input tensor to
`batch_flatten` has an unknown number of dimensions or unknown shape values,
the struct info inference for `batch_flatten` will correctly produce a
`TensorStructInfo` without a concrete shape. In this scenario,
`call.struct_info.shape` will be `None`, and attempting to access
`call.struct_info.shape.values` will result in an `AttributeError`, crashing
the compiler.
To make the legalization more robust, you should add a check to ensure
`call.struct_info.shape` is defined before attempting to use it. If it's not
defined, the operator should not be legalized, and the original call should be
returned. This is a common pattern in this file for other operators.
It would also be great to add a new test case in
`tests/python/relax/test_transform_legalize_ops_nn.py` to cover this scenario,
for example with an input tensor of `R.Tensor(ndim=4, dtype="float32")`, and
assert that `batch_flatten` is not legalized in this case.
```suggestion
def _nn_batch_flatten(bb: BlockBuilder, call: Call) -> Expr:
if call.struct_info.shape is None:
return call
return bb.call_te(topi.reshape, call.args[0],
call.struct_info.shape.values)
```
--
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]
