Searched refs:collectives_creator (Results 1 – 2 of 2) sorted by relevance
/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
D | spmd_partitioner.h | 213 const SPMDCollectiveOpsCreator& collectives_creator); 220 const SPMDCollectiveOpsCreator& collectives_creator, 235 const SPMDCollectiveOpsCreator& collectives_creator, bool per_dim_ag); 239 const SPMDCollectiveOpsCreator& collectives_creator,
|
D | spmd_partitioner.cc | 3502 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()
|