Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.cc1331 int64 target_i_size = target.tile_assignment().dim(i); in GetReshardAllToAllSourceTargetDims() local
1332 if (target_i_size == source_size) { in GetReshardAllToAllSourceTargetDims()
1337 auto j_it = source_size_to_dim[target_i_size].begin(); in GetReshardAllToAllSourceTargetDims()
1344 if (++j_it == source_size_to_dim[target_i_size].end()) { in GetReshardAllToAllSourceTargetDims()
1349 } else if (target_i_size % source_size == 0) { in GetReshardAllToAllSourceTargetDims()
1353 if (++j_it == source_size_to_dim[target_i_size].end()) { in GetReshardAllToAllSourceTargetDims()
1362 remove_entry(target_i_size, i, target_size_to_dim); in GetReshardAllToAllSourceTargetDims()
1364 remove_entry(target_i_size, j, source_size_to_dim); in GetReshardAllToAllSourceTargetDims()