[Relax][PyTorch] support for index.Tensor #17836
Conversation
Remaining TODOs Make sure that also works if we input a list of indices Write tests at the py and cpp level Cleanup
|
@tvm-bot rerun |
1 similar comment
|
@tvm-bot rerun |
|
@MasterJH5574 @tlopex this is ready for review. @tqchen note that this allows OpenAI's Whisper to be ingested by TVM using the exported program translator (after a few minor edits, including remove in-place operations). Correctness is achieved, as tested here https://github.com/hugolatendresse/catalyst-fleet/blob/a4a4be6/whisper/whisper_inlined.py The tvm-bot refuses to rerun the tests due to Ubuntu 20.04 LTS being retired - maybe this needs to be merged first |
hugolatendresse
changed the title
MasterJH5574
left a comment
MasterJH5574
left a comment
There was a problem hiding this comment.
Looks very great, nice done! I only left one comment for potentially better clarity.
| while isinstance(t, Var): | ||
| binding = bb.lookup_binding(t) | ||
| if not isinstance(binding, (Tuple, Var)): | ||
| break | ||
| t = binding | ||
|
|
||
| assert isinstance(t, (Tuple, Var)) |
There was a problem hiding this comment.
It might help understanding if we can briefly explain why backtracing t is needed. My understanding is that without backtracing it's still correct, but may result in trivial TupleGetItems that can be easily removed. Correct me if I'm wrong.
There was a problem hiding this comment.
@MasterJH5574 you are right that it is not needed. It would simply remove unneeded TupleGetItems like you said. I removed the backtracking and all tests still pass. Let me know if you prefer it like that or if I should revert (I'm indifferent).
MasterJH5574
left a comment
MasterJH5574
left a comment
There was a problem hiding this comment.
Looks great, thank you @hugolatendresse for landing this!
New op for advanced indexing + unit tests
2 participants
New op for advanced indexing + unit tests