Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.cc1139 auto temp_target = target.ReplicateOnLastTileDim() in ReshardWithAllToAll() local
1146 temp_target.tile_assignment().dim(target_dim))); in ReshardWithAllToAll()
1151 temp_target.tile_assignment().num_elements() / group_size); in ReshardWithAllToAll()
1152 temp_target.tile_assignment().Each( in ReshardWithAllToAll()
1157 group_id *= temp_target.tile_assignment().dim(dim) / group_size; in ReshardWithAllToAll()
1160 group_id *= temp_target.tile_assignment().dim(dim); in ReshardWithAllToAll()
1214 const Shape result_shape = MakePartitionedShape(base_shape_, temp_target); in ReshardWithAllToAll()
1220 result->set_sharding(temp_target); in ReshardWithAllToAll()