gemini-code-assist[bot] commented on code in PR #19671:
URL: https://github.com/apache/tvm/pull/19671#discussion_r3354039441
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -4203,6 +4203,17 @@ def _argreduce_select_last_index(bb, data, axis,
keepdims, op):
return relax.op.subtract(offset, flipped_idx)
+def _argreduce_sanitize_nan(bb, data, *, for_min):
+ """Match ONNX Runtime ArgMax/ArgMin behavior by making NaN win
comparisons."""
+ dtype = data.struct_info.dtype
+ if not _relax_dtype_is_floating_point(dtype):
+ return data
+ replacement = -_np.inf if for_min else _np.inf
+ return bb.emit(
+ relax.op.where(relax.op.isnan(data), relax.const(replacement, dtype),
data)
+ )
Review Comment:
### Optimization & Defensive Programming Opportunity
1. **Defensive Check**: Added a check `if not hasattr(data.struct_info,
"dtype")` to prevent potential `AttributeError` if `data.struct_info` is `None`
or does not have a `dtype` attribute.
2. **Constant Folding**: If `data` is a `relax.Constant`, we can perform the
NaN replacement directly at compile-time using NumPy. This avoids emitting
unnecessary `isnan`, `where`, and `const` operators in the Relax graph, keeping
the imported IR cleaner and more efficient.
```python
def _argreduce_sanitize_nan(bb, data, *, for_min):
"""Match ONNX Runtime ArgMax/ArgMin behavior by making NaN win
comparisons."""
if not hasattr(data.struct_info, "dtype"):
return data
dtype = data.struct_info.dtype
if not _relax_dtype_is_floating_point(dtype):
return data
replacement = -_np.inf if for_min else _np.inf
if isinstance(data, relax.Constant):
np_data = data.data.numpy()
if _np.any(_np.isnan(np_data)):
np_data = _np.where(_np.isnan(np_data), replacement, np_data)
return relax.const(np_data, dtype)
return data
return bb.emit(
relax.op.where(relax.op.isnan(data), relax.const(replacement,
dtype), data)
)
```
--
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]
