[Relax][ONNX][Transform] Add mode choice, new mode, and warning for take()
Summary
I was trying to resolve https://github.com/apache/tvm/issues/18004, where an ONNX model causes a segmentation fault in TVM but not in onnxruntime.
Why the seg fault occurs
This occurs because take()(used in Compress)defaults to fast mode, which deliberately segfaults on out-of-bounds indices. (I guess for the sake of fast speed)
Solution
onnxruntime's Compress ignores out-of-bounds indices, so I initially considered matching that behavior in TVM. However, the ONNX specification doesn’t actually mandate how to handle out-of-bounds indices
Question for maintainers
Would you prefer modifying Compress to slice out invalid indices in order to mirror onnxruntime? although I guess this is unnecessary.
Instead of resolving this seg-fault, for now, this PR focuses just on enhancing take().
Updates
- Add a
modeparameter totake()in Relax-
There are already several modes implemented in the transform layer, but Relax’s
take()defaults tofastwith no choice. I've exposed all modes so future use cases can pick as needed down the road.
-
- Add
NaNmode totake()-
I was initially adding a
raisemode(default ofnumpy.take()), but since there are no precedents raising runtime errors in the transform layer, I added anNaNmode that returnsNaNfor any out-of-bounds index. I'll remove this if this is unnecessary.
-
- Add unit tests covering all
take()modes-
The tests for take() modes had been lost during a refactor.
-
- Add a warning log for
fastmode along an axis-
take() along a flattened input shows a warning that lets users know about the potential seg fault, while
take()along an axis does not. I think bothtake()need to warn about it.
-
- Unify default modes in lower layers to
fastfor consistency with Relax-
The default mode for
take()in Relax is currentlyfast, I’ve made the lower layers usefastmode as well for consistency.
-
cc @tlopex
Sorry for the late response. Please rebase the code, and we can get it in
@Hzfengsy I've rebased the code according to the newly introduced reflection mechanism.