Searched refs:data_sharding (Results 1 – 2 of 2) sorted by relevance
639 HloSharding ScatterIndexSharding(const HloSharding& data_sharding, in ScatterIndexSharding() argument641 if (data_sharding.IsTileMaximal()) { in ScatterIndexSharding()642 return data_sharding; in ScatterIndexSharding()650 data_sharding.tile_assignment().dim(i)); in ScatterIndexSharding()656 if (data_sharding.ReplicateOnLastTileDim()) { in ScatterIndexSharding()658 data_sharding.tile_assignment().dimensions().back()); in ScatterIndexSharding()660 Array<int64> new_tile_assignment = data_sharding.tile_assignment(); in ScatterIndexSharding()663 return HloSharding::Replicate(data_sharding.metadata()); in ScatterIndexSharding()666 return data_sharding.ReplicateOnLastTileDim() in ScatterIndexSharding()668 data_sharding.metadata()) in ScatterIndexSharding()[all …]
128 HloSharding ScatterIndexSharding(const HloSharding& data_sharding,147 HloSharding ScatterEffectiveDataSharding(const HloSharding& data_sharding,