[stablehlo] Add batching dims to stablehlo.gather and stablehlo.scatter
#2259
Conversation
This file contains 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
tomnatan30
changed the title
stablehlo.gather and stablehlo.scatter
…lo.scatter" This reverts commit 39da3ab.
tomnatan30
commented
Apr 30, 2024
tomnatan30
commented
Apr 30, 2024
ghpvnist
requested changes
May 2, 2024
There was a problem hiding this comment.
Thanks for the thorough implementation! This is round 1 of X of my reviews :)
I haven't reviewed ops_stablehlo.mlir and verify_scatter.mlir because my brain is fried after reviewing this, and I expect the test names to change with the constraint number bump so I'll leave the review until the remaining feedbacks are addressed.
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 16, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to google-deepmind/tf2jax
that referenced
this pull request
Sep 21, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 21, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Sep 21, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Sep 21, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to google-deepmind/tf2jax
that referenced
this pull request
Sep 23, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 23, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Sep 23, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to google-deepmind/tf2jax
that referenced
this pull request
Sep 24, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 24, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Sep 24, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to google-deepmind/tf2jax
that referenced
this pull request
Sep 24, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 24, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to google-deepmind/tf2jax
that referenced
this pull request
Sep 24, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 24, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 24, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 24, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 24, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 25, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 25, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 25, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 25, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 25, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 25, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 25, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 25, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 25, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 25, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 647647825
copybara-service bot
pushed a commit
to jax-ml/jax
that referenced
this pull request
Sep 25, 2024
…attributes in stablehlo. This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota. See openxla/stablehlo#2259 PiperOrigin-RevId: 678649138
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.
Add
operand_batching_dimsandstart_indices_batching_dimsattributes tostablehlo.gather.operand_batching_dimsrefers to the dimensions of theoperandthat are treated as batch.start_indices_batching_dimsrefers to the dimensions of thestart_indicesthat are treated as batch. The corresponding dimension sizes must be equal. The semantics is equivalent to concatenating the outputs of the gather with each slices ofoperandandstart_indices.Similarly, add
input_batching_dimsandscatter_indices_batching_dimsattributes tostablehlo.scatter.input_batching_dimsrefers to the dimensions of each tensor ininputsthat are treated as batch.scatter_indices_batching_dimsrefers to the dimensions of thescatter_indicesthat are treated as batch.See #2084 for more information