Home
last modified time | relevance | path

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

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dspmd_partitioner_util.cc1540 if (creator.create_cross_partition_all_gather) { in GetPerGroupCollectiveOpsCreator()
1541 result.create_cross_partition_all_gather = in GetPerGroupCollectiveOpsCreator()
1546 return creator.create_cross_partition_all_gather( in GetPerGroupCollectiveOpsCreator()
H A Dspmd_partitioner.h156 create_cross_partition_all_gather; member
H A Dspmd_partitioner.cc1239 if (state_.collective_ops_creator.create_cross_partition_all_gather) { in ReplicatePartial()
4355 result = collectives_creator.create_cross_partition_all_gather( in AllGatherShardsInternal()
4373 result = collectives_creator.create_cross_partition_all_gather( in AllGatherShardsInternal()
H A Dspmd_partitioner_test.cc64 collective_ops_creator.create_cross_partition_all_gather = nullptr; in PartitionComputation()