Searched refs:index_sharding (Results 1 – 2 of 2) sorted by relevance
499 HloSharding GatherOutputSharding(const HloSharding& index_sharding, in GatherOutputSharding() argument501 if (index_sharding.IsTileMaximal()) { in GatherOutputSharding()502 return index_sharding; in GatherOutputSharding()514 index_sharding.tile_assignment().dim(new_tile_dimension)); in GatherOutputSharding()519 if (index_sharding.ReplicateOnLastTileDim()) { in GatherOutputSharding()521 index_sharding.tile_assignment().dimensions().back()); in GatherOutputSharding()524 Array<int64> new_tile_assignment = index_sharding.tile_assignment(); in GatherOutputSharding()527 return HloSharding::Replicate(index_sharding.metadata()); in GatherOutputSharding()530 return index_sharding.ReplicateOnLastTileDim() in GatherOutputSharding()532 index_sharding.metadata()) in GatherOutputSharding()[all …]
112 HloSharding GatherOutputSharding(const HloSharding& index_sharding,133 HloSharding ScatterDataSharding(const HloSharding& index_sharding,140 HloSharding ScatterEffectiveIndexSharding(const HloSharding& index_sharding,