gemini-code-assist[bot] commented on code in PR #19898:
URL: https://github.com/apache/tvm/pull/19898#discussion_r3487953832
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -1506,19 +1517,29 @@ def _impl_v14(cls, bb, inputs, attr, params):
x = inputs[0]
k = inputs[1] if len(inputs) > 1 else 0
- if len(inputs) > 1:
- k = get_constant(inputs[1], params)
- if isinstance(k, relax.Constant):
- k = int(k.data.numpy().item())
- else:
- raise ValueError("Currently only support constant k for Trilu
op.")
- else:
- k = 0
+ if isinstance(k, relax.Constant):
+ k = int(k.data.numpy().item())
+ if isinstance(k, int):
+ if upper:
+ return relax.op.triu(x, k)
+ return relax.op.tril(x, k)
+ # Dynamic k: build the mask explicitly so it works with any scalar k.
+ shape = x.ty.shape
+ m, n = shape[-2], shape[-1]
+ row_idx = relax.op.reshape(relax.op.arange(0, m, dtype="int64"), (m,
1))
+ col_idx = relax.op.reshape(relax.op.arange(0, n, dtype="int64"), (1,
n))
+ diff = relax.op.subtract(
+ relax.op.broadcast_to(col_idx, (m, n)),
+ relax.op.broadcast_to(row_idx, (m, n)),
+ )
Review Comment:
The explicit `relax.op.broadcast_to` calls on `col_idx` and `row_idx` are
redundant. In TVM Relax, element-wise binary operators like `relax.op.subtract`
support implicit broadcasting natively. Removing these explicit broadcasts
simplifies the generated Relax graph and avoids creating unnecessary
intermediate operators.
```python
diff = relax.op.subtract(col_idx, row_idx)
```
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -1170,6 +1170,17 @@ def _impl_v13(cls, bb, inputs, attr, params):
return relax.op.astype(inputs[0], to_type)
+class CastLike(OnnxOpConverter):
+ """Convert an onnx CastLike node into an equivalent Relax expression."""
+
+ @classmethod
+ def _impl_v15(cls, bb, inputs, attr, params):
+ data = inputs[0]
+ target = inputs[1]
+ target_dtype = target.ty.dtype.dtype
Review Comment:
Accessing `target.ty.dtype.dtype` directly assumes that `target.ty` is
populated and is a `TensorType`. If `target` is a `relax.Var` without type
annotation or has a different type structure, this will raise an
`AttributeError`. Consider using a safer fallback like `getattr` to handle
potentially missing type information, similar to how it is done in the `Cast`
converter.
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -1506,19 +1517,29 @@ def _impl_v14(cls, bb, inputs, attr, params):
x = inputs[0]
k = inputs[1] if len(inputs) > 1 else 0
- if len(inputs) > 1:
- k = get_constant(inputs[1], params)
- if isinstance(k, relax.Constant):
- k = int(k.data.numpy().item())
- else:
- raise ValueError("Currently only support constant k for Trilu
op.")
- else:
- k = 0
+ if isinstance(k, relax.Constant):
+ k = int(k.data.numpy().item())
+ if isinstance(k, int):
+ if upper:
+ return relax.op.triu(x, k)
+ return relax.op.tril(x, k)
+ # Dynamic k: build the mask explicitly so it works with any scalar k.
+ shape = x.ty.shape
+ m, n = shape[-2], shape[-1]
+ row_idx = relax.op.reshape(relax.op.arange(0, m, dtype="int64"), (m,
1))
+ col_idx = relax.op.reshape(relax.op.arange(0, n, dtype="int64"), (1,
n))
+ diff = relax.op.subtract(
+ relax.op.broadcast_to(col_idx, (m, n)),
+ relax.op.broadcast_to(row_idx, (m, n)),
+ )
+ k_int64 = relax.op.astype(k, "int64")
if upper:
- return relax.op.triu(x, k)
+ mask = relax.op.greater_equal(diff, k_int64)
else:
- return relax.op.tril(x, k)
+ mask = relax.op.less_equal(diff, k_int64)
+ mask = relax.op.broadcast_to(mask, shape)
+ return relax.op.where(mask, x, relax.const(0, x.ty.dtype))
Review Comment:
To be consistent with the rest of the file (e.g., lines 372, 383), you
should use `x.ty.dtype.dtype` to retrieve the string representation of the data
type for `relax.const`, rather than passing the `DataType` object `x.ty.dtype`
directly.
```suggestion
return relax.op.where(mask, x, relax.const(0, x.ty.dtype.dtype))
```
--
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]
