Searched refs:new_index_shard (Results 1 – 1 of 1) sorted by relevance
1840 auto new_index_shard = indices.sharding(); in GatherOperandsShardedAcrossParallelDims() local1842 int idx_parallel_tiles_num = new_index_shard.NumTiles(indices_parallel_dims); in GatherOperandsShardedAcrossParallelDims()1853 if (new_index_shard.IsReplicated()) { in GatherOperandsShardedAcrossParallelDims()1862 new_index_shard, in GatherOperandsShardedAcrossParallelDims()1863 CreateMatchingShardingOnDims(operand.shape(), new_index_shard, in GatherOperandsShardedAcrossParallelDims()1873 HloSharding* target = &new_index_shard; in GatherOperandsShardedAcrossParallelDims()1922 new_index_shard.tile_assignment().dim( in GatherOperandsShardedAcrossParallelDims()1931 operand_parallel_dims, new_index_shard, in GatherOperandsShardedAcrossParallelDims()1933 return GatherParallelDimSharding{new_index_shard, new_operand_shard}; in GatherOperandsShardedAcrossParallelDims()