Support shape transpose in hlo_sharding_util::ReshapeSharding
.
#64967
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.
Support shape transpose in
hlo_sharding_util::ReshapeSharding
.Before this cl,
hlo_sharding_util::ReshapeSharding
can handle the cases where source and target shapes can be transformed to each other by merging and splitting dimension sizes. It returnsstd::nullopt
if transpose is needed between source and target shapes.This cl extracts the gcd(source_sharding_tile_size, target_shape) when
source_shape % source_sharding_tile_size == 0
in the major dimensions. An example is shown below.output_sharding is
{devices=[2,1,1,1,3]<=[6] last_tile_dim_replicate}
. Before this cl, the output_sharding is{replicated}
.