Searched refs:new_operand_shard (Results 1 – 1 of 1) sorted by relevance
1841 auto new_operand_shard = operand.sharding(); in GatherOperandsShardedAcrossParallelDims() local1843 int op_parallel_tiles_num = new_operand_shard.NumTiles(operand_parallel_dims); in GatherOperandsShardedAcrossParallelDims()1855 CreateMatchingShardingOnDims(indices.shape(), new_operand_shard, in GatherOperandsShardedAcrossParallelDims()1858 new_operand_shard}; in GatherOperandsShardedAcrossParallelDims()1860 if (new_operand_shard.IsReplicated()) { in GatherOperandsShardedAcrossParallelDims()1874 HloSharding* to_adjust = &new_operand_shard; in GatherOperandsShardedAcrossParallelDims()1919 new_operand_shard.tile_assignment().dimensions(); in GatherOperandsShardedAcrossParallelDims()1925 auto operand_shard_tiles = new_operand_shard.tile_assignment(); in GatherOperandsShardedAcrossParallelDims()1927 new_operand_shard = in GatherOperandsShardedAcrossParallelDims()1928 AlignShardingOnDims(new_operand_shard.ReplicateOnLastTileDim() in GatherOperandsShardedAcrossParallelDims()[all …]