[Pallas/Mosaic TPU] Allow non-leading and non-matching batch dimensions in dot_general.
#33265
+109
−46
Conversation
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
copybara-service
bot
force-pushed
the
test_831015967
branch
6 times, most recently
from
ad20fe1 to
81cf5e5
Compare
…ns in `dot_general`. The constraints on `lhs_batch_dims` and `rhs_batch_dims` for `dot_general` in Pallas/Mosaic on TPU are now relaxed. Batch dimensions do not have to be at the front of the shape, and the dimension indices used for batching on the LHS and RHS can be different (which requires to update the semantics of `output_dim_order` in `tpu.dot_dimension_numbers`). The remaining gap compared to JAX is the lack of support for multiple batch dimensions. PiperOrigin-RevId: 831015967
copybara-service
bot
force-pushed
the
test_831015967
branch
from
81cf5e5 to
5fa26d0
Compare
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.
[Pallas/Mosaic TPU] Allow non-leading and non-matching batch dimensions in
dot_general.The constraints on
lhs_batch_dimsandrhs_batch_dimsfordot_generalin Pallas/Mosaic on TPU are now relaxed. Batch dimensions do not have to be at the front of the shape, and the dimension indices used for batching on the LHS and RHS can be different (which requires to update the semantics ofoutput_dim_orderintpu.dot_dimension_numbers).The remaining gap compared to JAX is the lack of support for multiple batch dimensions.