| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | all_reduce_blueconnect_test.cc | 77 auto all_gather = AllOf(op::Shape("f32[16]"), op::AllGather(all_reduce), in TEST_F() local
166 auto all_gather = AllOf(op::Shape("(f32[16], f32[32])"), in TEST_F() local
|
| D | all_reduce_blueconnect.cc | 239 HloInstruction* all_gather = in TryDecomposeAllReduce() local
|
| /external/tensorflow/tensorflow/dtensor/mlir/utils/ |
| D | collective_lowering.cc | 442 mlir::LogicalResult LowerAllGatherOp(mlir::TF::DTensorAllGatherOp all_gather) { in LowerAllGatherOp()
778 module.walk([&](mlir::TF::DTensorAllGatherOp all_gather) { in runOnOperation()
782 for (mlir::TF::DTensorAllGatherOp all_gather : all_gathers) in runOnOperation() local
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | collective_ops.py | 159 def all_gather(t, function
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | cross_device_utils.py | 434 def all_gather( member in CollectiveReplicaLauncher
|
| D | tpu_strategy.py | 1726 def all_gather(self, value, axis, experimental_hints=None): member in _TPUReplicaContext
|
| D | distribute_lib.py | 3302 def all_gather(self, value, axis, options=None): member in ReplicaContext
|
| /external/tensorflow/tensorflow/python/kernel_tests/ |
| D | collective_ops_test.py | 53 def all_gather(t, group_size, group_key, instance_key, *args, **kwargs): member in CollectiveOpsV1
73 def all_gather(t, group_size, group_key, instance_key, *args, **kwargs): member in CollectiveOpsV2
|
| /external/tensorflow/tensorflow/dtensor/mlir/ |
| D | collectives.cc | 93 mlir::TF::DTensorAllGatherOp all_gather = in EmitAllGather() local
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/ |
| D | hlo_function_importer.cc | 973 auto all_gather = Cast<HloAllGatherInstruction>(instruction); in ImportInstructionImpl() local
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | mhlo_to_lhlo_with_xla.cc | 1161 auto* all_gather = xla::Cast<xla::HloAllGatherInstruction>(instr); in EmitAllGatherOp() local
|