Home
last modified time | relevance | path

Searched defs:all_to_all (Results 1 – 12 of 12) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dall_to_all_decomposer.cc39 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 Dhlo_matchers_test.cc337 std::unique_ptr<HloInstruction> all_to_all = in TEST_F() local
H A Dhlo_verifier.cc487 auto* all_to_all = Cast<HloAllToAllInstruction>(hlo); in HandleAllToAll() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dgpu_layout_assignment.cc391 auto* all_to_all = Cast<HloAllToAllInstruction>(instruction); in AddBackendConstraints() local
/aosp_15_r20/external/tensorflow/tensorflow/python/tpu/ops/
H A Dtpu_ops.py40 def all_to_all(x, function
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dfft_handler.cc165 auto all_to_all = collective_ops_creator.create_cross_partition_all_to_all( in ShuffleDataWithAllToAll() local
H A Dspmd_partitioner_test.cc1178 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 Dspmd_partitioner.cc1585 auto all_to_all = in ReshardWithAllToAll() local
/aosp_15_r20/external/pytorch/torch/distributed/nn/
H A Dfunctional.py158 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 Dhlo_function_importer.cc1008 auto all_to_all = Cast<HloAllToAllInstruction>(instruction); in ImportInstructionImpl() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
H A Dmhlo_to_lhlo_with_xla.cc1147 auto* all_to_all = xla::Cast<xla::HloAllToAllInstruction>(instr); in EmitAllToAllOp() local
/aosp_15_r20/external/pytorch/torch/distributed/
H A Ddistributed_c10d.py4007 def all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False): function