Searched defs:partition_subgroups (Results 1 – 3 of 3) sorted by relevance
1486 const std::vector<std::vector<int64_t>>& partition_subgroups) { in GetPerGroupCollectiveOpsCreator()1508 int64_t channel_id) { in GetPerGroupCollectiveOpsCreator()1544 const std::vector<std::vector<int64_t>>& partition_subgroups, in GetPerGroupCollectiveOpsCreator()1545 int64_t channel_id, int64_t all_gather_dimension) { in GetPerGroupCollectiveOpsCreator()
162 std::vector<int64_t> partition_subgroups(num_partitions); in ShuffleDataWithAllToAll() local
4351 auto partition_subgroups = in AllGatherShardsInternal() local4370 auto partition_subgroups = in AllGatherShardsInternal() local4443 auto partition_subgroups = in AllReduceAlongShardingDimsInternal() local4453 auto partition_subgroups = in AllReduceAlongShardingDimsInternal() local