Home
last modified time | relevance | path

Searched refs:new_tile_assignment_dims (Results 1 – 1 of 1) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.cc1882 auto new_tile_assignment_dims = to_adjust->tile_assignment().dimensions(); in GatherOperandsShardedAcrossParallelDims() local
1894 new_tile_assignment_dims.back() % ratio != 0) { in GatherOperandsShardedAcrossParallelDims()
1897 new_tile_assignment_dims[to_adjust_dims[i]] *= ratio; in GatherOperandsShardedAcrossParallelDims()
1898 new_tile_assignment_dims.back() /= ratio; in GatherOperandsShardedAcrossParallelDims()
1900 CHECK_GE(new_tile_assignment_dims.back(), 1); in GatherOperandsShardedAcrossParallelDims()
1902 if (new_tile_assignment_dims.back() == 1) { in GatherOperandsShardedAcrossParallelDims()
1903 new_tile_assignment_dims.pop_back(); in GatherOperandsShardedAcrossParallelDims()
1907 new_tile_assignment.Reshape(new_tile_assignment_dims); in GatherOperandsShardedAcrossParallelDims()