gemini-code-assist[bot] commented on code in PR #19979:
URL: https://github.com/apache/tvm/pull/19979#discussion_r3560488804
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -3576,6 +3576,11 @@ def _impl_v15(cls, bb, inputs, attr, params):
epsilon = attr.get("epsilon", 1e-05)
momentum = attr.get("momentum", 0.9)
training_mode = attr.get("training_mode", 0)
+ data_dtype = data.ty.dtype
+ scale, bias, mean, var = [
+ relax.op.astype(x, data_dtype) if x.ty.dtype != data_dtype else x
+ for x in [scale, bias, mean, var]
+ ]
Review Comment:
When the ONNX parameters (`scale`, `bias`, `mean`, `var`) are constants
(which is the default when `keep_params_in_input=False`), casting them using
`relax.op.astype` introduces unnecessary `astype` call nodes in the Relax AST.
We can optimize this by directly casting the constant's underlying numpy
array and recreating a `relax.Constant` with the target data type. This keeps
the imported Relax graph cleaner and avoids relying on subsequent constant
folding passes to clean up the cast operations.
```python
data_dtype = data.ty.dtype
scale, bias, mean, var = [
relax.const(x.data.numpy().astype(str(data_dtype)),
str(data_dtype))
if isinstance(x, relax.Constant) and x.ty.dtype != data_dtype
else relax.op.astype(x, data_dtype)
if x.ty.dtype != data_dtype
else x
for x in [scale, bias, mean, var]
]
```
--
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]
