Searched defs:device_groups (Results 1 – 3 of 3) sorted by relevance
307 std::vector<std::vector<int64>> device_groups; member
1398 std::vector<std::vector<int64>> device_groups(Product(group_dim_sizes)); in GroupShardingOnDims() local1574 const std::vector<std::vector<int64>>& device_groups, SpmdBuilder* b) { in GetInGroupPartitionId()1592 const std::vector<std::vector<int64>>& device_groups) { in GetPerGroupCollectiveOpsCreator()1673 const std::vector<std::vector<int64>>& device_groups, SpmdBuilder* b) { in CreatePerGroupPartitioningState()1695 const std::vector<std::vector<int64>>& device_groups, in PerGroupSliceFromReplicated()1755 const std::vector<std::vector<int64>>& device_groups) { in FindMatchingPartitionedDimsForGrouping()
3433 std::vector<ReplicaGroup> device_groups; in GetDefaultCollectiveOpsCreator() local3475 std::vector<ReplicaGroup> device_groups; in GetDefaultCollectiveOpsCreator() local