Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[XLA:SPMD] Fix ReshapeSharding for dimensions of size 1 and >1 partitions. #17787

Closed
wants to merge 0 commits into from

Conversation

copybara-service[bot]
Copy link

[XLA:SPMD] Fix ReshapeSharding for dimensions of size 1 and >1 partitions.

If there is a source dimension satisfying the following conditions, we replicate the source sharding along this dimension since the source sharding cannot be propagated along this dimension.

  1. its size is 1
  2. its partitions is greater than 1
  3. there is no corresponding target dimension

An example is shown below. Please refer to the added examples in hlo_sharding_util_test.cc.

input shape: [1,2,16]
input sharding: [3,2,2]<=[12]
output shape: [2,16]

output sharding before this cl: [2,2,3]<=[12] last_tile_dim_replicate
output sharding with this cl: [2,2,3]<=[3,2,2]T(1,2,0) last_tile_dim_replicate

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant