Repository navigation
Conversation
topi.index_put flattened the broadcast index shape and read values at each flat position, but never broadcast values to that shape. Values that needed broadcasting (for example x[:, [0, 2]] = v with v of shape (2, 1)) therefore indexed past the end of values and wrote garbage. Broadcast values to the indexed-region shape following NumPy rules before building the store, matching torch.index_put semantics for both the overwrite and accumulate paths. Add topi and Relax regression tests covering (2, 1), (2,) and scalar values, with accumulate disabled and enabled.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
topi.index_putbroadcasts the index tensors together to obtain the indexedregion shape, but it read
valuesat each flat position of that shape withoutbroadcasting
valuesfirst. Whenvaluesneeded broadcasting (for examplex[:, [0, 2]] = vwithvof shape(2, 1)), the generated kernel indexedpast the end of
valuesand produced wrong values. The Relaxindex_putlegalization passes
valuesthrough unchanged, so the bug was reachable fromRelax and from the PyTorch frontend.
Broadcast
valuesto the indexed-region shape with NumPy rules before buildingthe store, for both the overwrite and
accumulate=Truepaths. This matchestorch.index_putsemantics.Tests:
tests/python/topi/test_topi_ir_builder.py: newtest_index_put_broadcast_values, covering(2, 1),(2,)and scalarvalues with and without accumulation.
tests/python/relax/test_op_manipulate.py: newtest_index_put_broadcast_values, an end-to-end Relax compile/run test withindices_0of shape(2, 1)andvaluesof shape(2, 1).Local verification:
pytest tests/python/topi/test_topi_ir_builder.py tests/python/relax/test_op_manipulate.py -n 8-> 171 passed.
main'stopi.index_put(writes[[10, 1, 20], [0, 4, 0]]instead of[[10, 1, 10], [20, 4, 20]]).pre-commit run --files ...passes.Fixes #20925.