Searched refs:source_size_to_dim (Results 1 – 1 of 1) sorted by relevance
1294 std::map<int64, std::vector<int64>> source_size_to_dim; in GetReshardAllToAllSourceTargetDims() local1300 source_size_to_dim[source.tile_assignment().dim(i)].push_back(i); in GetReshardAllToAllSourceTargetDims()1305 if (source_size_to_dim.empty() || in GetReshardAllToAllSourceTargetDims()1306 source_size_to_dim.size() != target_size_to_dim.size()) { in GetReshardAllToAllSourceTargetDims()1309 for (const auto& entry : source_size_to_dim) { in GetReshardAllToAllSourceTargetDims()1328 while (!source_size_to_dim.empty()) { in GetReshardAllToAllSourceTargetDims()1329 int64 source_size = source_size_to_dim.begin()->first; in GetReshardAllToAllSourceTargetDims()1330 int64 i = source_size_to_dim.begin()->second.back(); in GetReshardAllToAllSourceTargetDims()1333 remove_entry(source_size, i, source_size_to_dim); in GetReshardAllToAllSourceTargetDims()1337 auto j_it = source_size_to_dim[target_i_size].begin(); in GetReshardAllToAllSourceTargetDims()[all …]