Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.h213 const SPMDCollectiveOpsCreator& collectives_creator);
220 const SPMDCollectiveOpsCreator& collectives_creator,
235 const SPMDCollectiveOpsCreator& collectives_creator, bool per_dim_ag);
239 const SPMDCollectiveOpsCreator& collectives_creator,
Dspmd_partitioner.cc3502 const SPMDCollectiveOpsCreator& collectives_creator) { in AllGatherShards() argument
3504 selected_dims, collectives_creator, in AllGatherShards()
3511 const SPMDCollectiveOpsCreator& collectives_creator, bool per_dim_ag) { in AllGatherShardsInternal() argument
3533 result = collectives_creator.create_cross_partition_all_gather( in AllGatherShardsInternal()
3543 result = collectives_creator.create_cross_partition_all_gather( in AllGatherShardsInternal()
3601 const SPMDCollectiveOpsCreator& collectives_creator, in AllReduceAlongShardingDims() argument
3604 b, operand, sharding, next_channel_id, selected_dims, collectives_creator, in AllReduceAlongShardingDims()
3611 const SPMDCollectiveOpsCreator& collectives_creator, in AllReduceAlongShardingDimsInternal() argument
3616 return collectives_creator.create_cross_partition_all_reduce( in AllReduceAlongShardingDimsInternal()
3626 result = collectives_creator.create_cross_partition_all_reduce( in AllReduceAlongShardingDimsInternal()