Searched defs:all_to_all (Results 1 – 12 of 12) sorted by relevance
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
H A D | all_to_all_decomposer.cc | 39 auto* all_to_all = DynCast<HloAllToAllInstruction>(instruction); in InstructionMatchesPattern() local 57 auto* all_to_all = Cast<HloAllToAllInstruction>(instruction); in ExpandInstruction() local 87 HloInstruction* all_to_all = in ExpandInstruction() local
|
H A D | hlo_matchers_test.cc | 337 std::unique_ptr<HloInstruction> all_to_all = in TEST_F() local
|
H A D | hlo_verifier.cc | 487 auto* all_to_all = Cast<HloAllToAllInstruction>(hlo); in HandleAllToAll() local
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
H A D | gpu_layout_assignment.cc | 391 auto* all_to_all = Cast<HloAllToAllInstruction>(instruction); in AddBackendConstraints() local
|
/aosp_15_r20/external/tensorflow/tensorflow/python/tpu/ops/ |
H A D | tpu_ops.py | 40 def all_to_all(x, function
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
H A D | fft_handler.cc | 165 auto all_to_all = collective_ops_creator.create_cross_partition_all_to_all( in ShuffleDataWithAllToAll() local
|
H A D | spmd_partitioner_test.cc | 1178 auto all_to_all = in TEST_F() local 1624 auto all_to_all = in TEST_F() local 1659 auto all_to_all = in TEST_F() local 7319 auto all_to_all = AllOf(op::Shape("f32[4,4,2,4,4]"), op::AllToAll(reshape)); in TEST_F() local 7343 auto all_to_all = op::AllToAll( in TEST_F() local 7366 auto all_to_all = op::AllToAll( in TEST_F() local 7636 auto all_to_all = AllOf(op::Shape("f32[2,2,4,12,32]"), op::AllToAll(reshape)); in TEST_F() local
|
H A D | spmd_partitioner.cc | 1585 auto all_to_all = in ReshardWithAllToAll() local
|
/aosp_15_r20/external/pytorch/torch/distributed/nn/ |
H A D | functional.py | 158 def all_to_all(output_tensor_list, input_tensor_list, group=group.WORLD): function
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/ |
H A D | hlo_function_importer.cc | 1008 auto all_to_all = Cast<HloAllToAllInstruction>(instruction); in ImportInstructionImpl() local
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
H A D | mhlo_to_lhlo_with_xla.cc | 1147 auto* all_to_all = xla::Cast<xla::HloAllToAllInstruction>(instr); in EmitAllToAllOp() local
|
/aosp_15_r20/external/pytorch/torch/distributed/ |
H A D | distributed_c10d.py | 4007 def all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False): function
|