gemini-code-assist[bot] commented on code in PR #19680:
URL: https://github.com/apache/tvm/pull/19680#discussion_r3367806383
##########
python/tvm/relax/transform/legalize_ops/linear_algebra.py:
##########
@@ -45,6 +45,8 @@ def te_matmul(a: te.Tensor, b: te.Tensor) -> te.Tensor:
b_relax = relax.Var("b", relax.TensorStructInfo(b.shape))
f_infer_sinfo = call.op.get_attr("FInferStructInfo")
output_shape = f_infer_sinfo(relax.op.matmul(a_relax, b_relax),
bb).shape
+ if isinstance(a_shape[-1], trix.IntImm) and a_shape[-1] == 0:
Review Comment:
There is a typo `trix` instead of `tirx`. This will cause a `NameError` at
runtime when legalizing matmul with a zero-size dimension.
```suggestion
if isinstance(a_shape[-1], tirx.IntImm) and a_shape[-1] == 0:
```
##########
python/tvm/relax/transform/legalize_ops/statistical.py:
##########
@@ -17,15 +17,54 @@
# pylint: disable=invalid-name
"""Default legalization function for statistical operators."""
+from typing import Callable
+
from tvm import te, tirx, topi
from ...block_builder import BlockBuilder
from ...expr import Call, Expr
from .common import LegalizeFunc, TEFunc, register_legalize
-def _statistical(te_func: TEFunc) -> LegalizeFunc:
+def _normalize_reduction_axes(axis: list[int] | None, ndim: int) -> list[int]:
+ if axis is None:
+ return list(range(ndim))
+
+ axes = []
+ for dim in axis:
+ if isinstance(dim, trix.IntImm):
Review Comment:
There is a typo `trix` instead of `tirx`. This will cause a `NameError` at
runtime when normalizing reduction axes.
```suggestion
if isinstance(dim, tirx.IntImm):
```
##########
python/tvm/relax/transform/legalize_ops/statistical.py:
##########
@@ -17,15 +17,54 @@
# pylint: disable=invalid-name
"""Default legalization function for statistical operators."""
+from typing import Callable
+
from tvm import te, tirx, topi
from ...block_builder import BlockBuilder
from ...expr import Call, Expr
from .common import LegalizeFunc, TEFunc, register_legalize
-def _statistical(te_func: TEFunc) -> LegalizeFunc:
+def _normalize_reduction_axes(axis: list[int] | None, ndim: int) -> list[int]:
+ if axis is None:
+ return list(range(ndim))
+
+ axes = []
+ for dim in axis:
+ if isinstance(dim, trix.IntImm):
+ dim = dim.value
+ dim = int(dim)
+ axes.append(dim + ndim if dim < 0 else dim)
+ return axes
+
+
+def _has_const_zero_reduction_dim(call: Call) -> bool:
+ input_shape = call.args[0].struct_info.shape
+ if not isinstance(input_shape, ShapeExpr):
+ return false
Review Comment:
Use Python's capitalized `False` instead of lowercase `false` to avoid a
`NameError`.
```suggestion
return False
```
##########
python/tvm/relax/transform/legalize_ops/statistical.py:
##########
@@ -17,15 +17,54 @@
# pylint: disable=invalid-name
"""Default legalization function for statistical operators."""
+from typing import Callable
+
from tvm import te, tirx, topi
from ...block_builder import BlockBuilder
from ...expr import Call, Expr
Review Comment:
The class `ShapeExpr` is used in `_has_const_zero_reduction_dim` but is not
imported in this file. Please import it from `...expr` to avoid a `NameError`.
```suggestion
from ...expr import Call, Expr, ShapeExpr
```
--
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]
