gemini-code-assist[bot] commented on code in PR #19975:
URL: https://github.com/apache/tvm/pull/19975#discussion_r3556841485
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -527,11 +527,37 @@ class Div(BinaryBase):
numpy_op = _np.divide
relax_op = relax.op.divide
+ @staticmethod
+ def _as_scalar_prim_expr(expr, dtype):
+ if tvm.ir.is_prim_expr(expr):
+ return expr
+ if isinstance(expr, relax.Constant) and expr.data.shape == ():
+ return tirx.const(expr.data.numpy().item(), dtype)
+ return None
Review Comment:
The static method `_as_scalar_prim_expr` checks if `expr` is a `PrimExpr`
using `tvm.ir.is_prim_expr(expr)`. However, shape-derived values in Relax are
often wrapped in `relax.PrimValue` (for example, the output of `Gather` on a
`ShapeExpr` returns `relax.prim_value(shape_val)`). Since `relax.PrimValue` is
a subclass of `relax.Expr` and not `tvm.ir.PrimExpr`, `tvm.ir.is_prim_expr`
will return `False` for it. We should also check for and unwrap
`relax.PrimValue` to extract its underlying `PrimExpr` via `.value`.
```suggestion
@staticmethod
def _as_scalar_prim_expr(expr, dtype):
if isinstance(expr, relax.PrimValue):
return expr.value
if tvm.ir.is_prim_expr(expr):
return expr
if isinstance(expr, relax.Constant) and expr.data.shape == ():
return tirx.const(expr.data.numpy().item(), dtype)
return None
```
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -540,9 +566,15 @@ def _impl_v7(cls, bb, inputs, attr, params):
if not (lhs_is_integer and rhs_is_integer):
return cls.base_impl(bb, inputs, attr, params)
- if isinstance(inputs[1], relax.Constant) and
bool(_np.any(inputs[1].data.numpy() == 0)):
+ if cls._is_zero(inputs[1]):
raise ValueError("ONNX Div with integer inputs encountered divisor
value 0.")
+ has_prim_expr = any(tvm.ir.is_prim_expr(inp) for inp in inputs)
Review Comment:
Similar to `_as_scalar_prim_expr`, `has_prim_expr` is computed using
`any(tvm.ir.is_prim_expr(inp) for inp in inputs)`. If the inputs contain a
`relax.PrimValue` (which is a Relax expression, not a TIR `PrimExpr`),
`tvm.ir.is_prim_expr` will return `False`. This will cause `has_prim_expr` to
be `False`, bypassing the `tirx.truncdiv` path and falling back to `base_impl`,
which can fail with an `AttributeError` because `relax.PrimValue` does not have
a `data` attribute. We should also check for `relax.PrimValue` when determining
`has_prim_expr`.
```suggestion
has_prim_expr = any(tvm.ir.is_prim_expr(inp) or isinstance(inp,
relax.PrimValue) for inp in inputs)
```
--
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]
