/external/tensorflow/tensorflow/compiler/xla/service/ |
D | all_to_all_decomposer.cc | 39 auto* all_to_all = DynCast<HloAllToAllInstruction>(instruction); in InstructionMatchesPattern() local 40 if (all_to_all == nullptr) { in InstructionMatchesPattern() 44 if (all_to_all->constrain_layout()) { in InstructionMatchesPattern() 47 if (all_to_all->shape().IsTuple()) { in InstructionMatchesPattern() 53 return all_to_all->shape().rank() < min_array_rank_; in InstructionMatchesPattern() 57 auto* all_to_all = Cast<HloAllToAllInstruction>(instruction); in ExpandInstruction() local 58 int64 split_dim = *all_to_all->split_dimension(); in ExpandInstruction() 60 all_to_all->replica_groups().empty() in ExpandInstruction() 62 : all_to_all->replica_groups()[0].replica_ids_size(); in ExpandInstruction() 64 all_to_all->shape().dimensions(split_dim) / all_to_all_group_size; in ExpandInstruction() [all …]
|
D | hlo_verifier.cc | 292 auto* all_to_all = Cast<HloAllToAllInstruction>(hlo); in HandleAllToAll() local 295 TF_RET_CHECK(all_to_all != nullptr); in HandleAllToAll() 296 if (all_to_all->split_dimension()) { in HandleAllToAll() 319 if (all_to_all->split_dimension()) { in HandleAllToAll() 323 hlo->operand(0)->shape(), *all_to_all->split_dimension(), in HandleAllToAll() 324 *all_to_all->split_dimension(), split_count)); in HandleAllToAll()
|
/external/tensorflow/tensorflow/python/tpu/ops/ |
D | tpu_ops.py | 45 def all_to_all(x, function 69 return gen_tpu_ops.all_to_all( 84 gen_tpu_ops.all_to_all(
|
/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
D | fft_handler.cc | 167 auto all_to_all = collective_ops_creator.create_cross_partition_all_to_all( in ShuffleDataWithAllToAll() local 169 return all_to_all; in ShuffleDataWithAllToAll()
|
D | spmd_partitioner_test.cc | 1144 auto all_to_all = in TEST_F() local 1146 auto reshard_lhs = AllOf(op::Reshape(op::Transpose(all_to_all)), in TEST_F() 1586 auto all_to_all = in TEST_F() local 1588 auto reshard = AllOf(op::Reshape(op::Transpose(all_to_all))); in TEST_F() 1621 auto all_to_all = in TEST_F() local 1623 auto reshard = op::Reshape(op::Transpose(all_to_all)); in TEST_F() 4939 auto all_to_all = AllOf(op::Shape("f32[4,4,2,4,4]"), op::AllToAll(reshape)); in TEST_F() local 4940 auto xpose = AllOf(op::Shape("f32[2,4,4,4,4]"), op::Transpose(all_to_all)); in TEST_F() 4963 auto all_to_all = op::AllToAll( in TEST_F() local 4966 AllOf(op::Shape("f32[2,4]"), op::Reshape(op::Transpose(all_to_all))); in TEST_F() [all …]
|
D | spmd_partitioner.cc | 1184 auto all_to_all = in ReshardWithAllToAll() local 1193 for (int64 i = 0; i < all_to_all->shape().rank(); ++i) { in ReshardWithAllToAll() 1203 ShapeInference::InferTransposeShape(all_to_all->shape(), permutation) in ReshardWithAllToAll() 1205 all_to_all, permutation)); in ReshardWithAllToAll()
|
/external/tensorflow/tensorflow/python/distribute/ |
D | tpu_strategy.py | 1324 result = tpu_ops.all_to_all( 1565 unordered_output = tpu_ops.all_to_all( 1578 replica_ids = tpu_ops.all_to_all(
|
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
D | mhlo_to_lhlo_with_xla.cc | 1051 auto* all_to_all = xla::Cast<xla::HloAllToAllInstruction>(instr); in EmitAllToAllOp() local 1054 if (all_to_all->split_dimension().has_value()) { in EmitAllToAllOp() 1056 builder_.getI64IntegerAttr(*all_to_all->split_dimension())); in EmitAllToAllOp()
|
/external/tensorflow/tensorflow/compiler/xla/client/ |
D | xla_builder.cc | 2894 XlaOp all_to_all, in AllToAllArray() 2897 return all_to_all; in AllToAllArray() 2908 all_to_all = Reshape(all_to_all, sizes); in AllToAllArray() 2918 all_to_all = Transpose(all_to_all, permutation); in AllToAllArray() 2919 return Reshape(all_to_all_shape, all_to_all); in AllToAllArray()
|
/external/tensorflow/tensorflow/compiler/mlir/hlo/tests/ |
D | lhlo_ops.mlir | 48 "lmhlo.all_to_all"(%input0, %output) 59 "lmhlo.all_to_all"(%input0, %output)
|
D | ops.mlir | 18 %0 = "mhlo.all_to_all"(%data) { 31 %0 = "mhlo.all_to_all"(%data) { 44 %0 = "mhlo.all_to_all"(%data) {
|
/external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/ |
D | lhlo_ops.td | 572 def LHLO_AllToAllOp : LHLO_CollectiveCommunicationOp<"all_to_all">,
|
D | hlo_ops.td | 599 def HLO_AllToAllOp : HLO_Op<"all_to_all",
|
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/hlo_to_lhlo_with_xla/ |
D | hlo_text_to_lhlo_no_opt.hlotxt | 619 // CHECK: "lmhlo.all_to_all"
|
/external/tensorflow/tensorflow/python/tpu/ |
D | tensor_tracer.py | 1574 return tpu_ops.all_to_all(
|
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/ |
D | legalize-tf.mlir | 4885 // CHECK: mhlo.all_to_all
|