Home
last modified time | relevance | path

Searched refs:AllToAll (Results 1 – 25 of 37) sorted by relevance

12

/external/tensorflow/tensorflow/core/common_runtime/
Dall_to_all.cc42 AllToAll::AllToAll() in AllToAll() function in tensorflow::AllToAll
45 StatusCallback AllToAll::CheckCounterAndCallDone() { in CheckCounterAndCallDone()
78 Status AllToAll::InitializeCollectiveContext( in InitializeCollectiveContext()
94 void AllToAll::Run(StatusCallback done) { in Run()
123 void AllToAll::DispatchSend(int src_rank, int target_rank, const Tensor* tensor, in DispatchSend()
135 void AllToAll::DispatchRecv(int src_rank, int target_rank, Tensor* tensor, in DispatchRecv()
149 REGISTER_COLLECTIVE(AllToAll, AllToAll);
Dall_to_all.h30 class AllToAll : public CollectiveImplementationInterface {
32 AllToAll();
/external/tensorflow/tensorflow/security/advisory/
Dtfsa-2021-176.md1 ## TFSA-2021-176: Integer division by 0 in `tf.raw_ops.AllToAll`
7 The [shape inference code for `AllToAll`](https://github.com/tensorflow/tensorflow/blob/8d72537c6ab…
14 return tf.raw_ops.AllToAll(
/external/tensorflow/tensorflow/core/ops/
Dtpu_cross_replica_ops_test.cc24 TEST(AllToAll, UnknownRank) { in TEST() argument
31 TEST(AllToAll, KnownRankUnknownDims) { in TEST() argument
/external/tensorflow/tensorflow/core/ops/compat/ops_history_v2/
DAllToAll.pbtxt2 name: "AllToAll"
39 name: "AllToAll"
92 name: "AllToAll"
/external/tensorflow/tensorflow/core/ops/compat/ops_history_v1/
DAllToAll.pbtxt2 name: "AllToAll"
39 name: "AllToAll"
/external/tensorflow/tensorflow/core/api_def/base_api/
Dapi_def_AllToAll.pbtxt2 graph_op_name: "AllToAll"
/external/tensorflow/tensorflow/core/tpu/kernels/
Dcross_replica_ops.cc84 0, xla::AllToAll(ctx->Input(0), split_dimension_, concat_dimension_, in Compile()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_test.cc245 AllOf(op::Copy(op::Reshape(op::Transpose(op::AllToAll(AllOf( in TEST_F()
265 AllOf(op::Copy(op::Slice(op::Reshape(AllOf(op::Transpose(op::AllToAll( in TEST_F()
1179 AllOf(op::AllToAll(op::Reshape(lhs)), op::Shape("f32[64,2,112,224,3]")); in TEST_F()
1428 op::AllToAll(op::Reshape(op::Pad(operand, op::Constant()))))); in TEST_F()
1625 AllOf(op::AllToAll(op::Reshape(lhs)), op::Shape("f32[2,64,28,56,64]")); in TEST_F()
1660 AllOf(op::AllToAll(op::Reshape(rhs)), op::Shape("f32[64,2,14,28,64]")); in TEST_F()
3038 auto resahrd = AllOf(op::Reshape(op::Transpose(op::Reshape(op::AllToAll()))), in TEST_F()
3086 auto resahrd = AllOf(op::Reshape(op::Transpose(op::Reshape(op::AllToAll()))), in TEST_F()
3155 op::Reshape(op::Transpose(op::AllToAll(op::Reshape(op::Parameter(0))))); in TEST_F()
3179 op::AllToAll(op::Reshape(local_reshape)))))); in TEST_F()
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Djitrt_custom_calls.cc1922 struct AllToAll { struct
1930 static AllToAll Handler() { return AllToAll(); } in Handler() argument
1934 Error AllToAll::operator()(const ServiceExecutableRunOptions* run_options, in operator ()()
1967 static bool AllToAll(runtime::KernelContext* ctx, void** args, void** attrs) { in AllToAll() function
1979 .To<RuntimeChecks()>(AllToAll::Handler()) in AllToAll()
2203 lib.Insert("xla.gpu.all_to_all", &xla::gpu::AllToAll); in JitRtGpuCustomCalls()
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/include/mlir-hlo/Dialect/mhlo/IR/
Dhlo_ops_base_attrs.td113 // AllToAll). Non-positive channel_id handle is equivalent to no channel id.
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dsimple_orc_jit.cc278 REGISTER_CPU_RUNTIME_SYMBOL(AllToAll); in RegisterKnownJITSymbols()
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/stablehlo/dialect/
DStablehloAttrs.td117 // AllToAll). Non-positive channel_id handle is equivalent to no channel id.
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder_test.cc436 TEST_F(XlaBuilderTest, AllToAll) { in TEST_F() argument
439 AllToAll(x, /*split_dimension=*/1, /*concat_dimension=*/0, in TEST_F()
456 AllToAll(x, /*split_dimension=*/0, /*concat_dimension=*/0, in TEST_F()
Dxla_builder.h769 XlaOp AllToAll(XlaOp operand, int64_t split_dimension,
1381 friend XlaOp AllToAll(XlaOp operand, int64_t split_dimension,
2368 XlaOp AllToAll(XlaOp operand, int64_t split_dimension, int64_t concat_dimension,
Dxla_builder.cc3077 XlaOp XlaBuilder::AllToAll(XlaOp operand, int64_t split_dimension, in AllToAll() function in xla::XlaBuilder
4690 XlaOp AllToAll(const XlaOp operand, int64_t split_dimension, in AllToAll() function
4694 return operand.builder()->AllToAll(operand, split_dimension, concat_dimension, in AllToAll()
/external/tensorflow/tensorflow/compiler/xla/tools/
Dhlo_control_flow_flattening_test.cc395 TEST_F(HloControlFlowFlatteningTest, AllToAll) { in TEST_F() argument
/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_matchers.h245 HLO_MATCHER(AllToAll);
Dhlo_parser_test.cc1836 "AllToAll", in CreateTestCases()
1837 R"(HloModule AllToAll, entry_computation_layout={(f32[128,32]{0,1})->(f32[128,32]{0,1})} in CreateTestCases()
1839 ENTRY AllToAll { in CreateTestCases()
Dhlo.proto186 // AllToAll). Non-positive channel_id is equivalent to no channel id.
/external/tensorflow/tensorflow/compiler/xla/python/xla_extension/
Dops.pyi114 def AllToAll(
/external/tensorflow/tensorflow/compiler/xla/python/
Dops.cc95 ops.def("AllToAll", &AllToAll, py::arg("operand"), py::arg("split_dimension"), in BuildOpsSubmodule()
/external/tensorflow/tensorflow/compiler/xla/g3doc/
Doperation_semantics.md115 ## AllToAll section in Operation Semantics
118 [`XlaBuilder::AllToAll`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder…
120 AllToAll is a collective operation that sends data from all cores to all cores.
133 retrieved using [`ReplicaId`](#replicaid)). AllToAll will be applied within
135 {4,5,0}}` means that an AllToAll will be applied within replicas `{1, 2,
137 in the same order of 1, 2, 3. Then, another AllToAll will be applied within
148 <b> `AllToAll(operand, split_dimension, concat_dimension, split_count,
178 AllToAll(x, /*split_dimension=*/1, /*concat_dimension=*/0, /*split_count=*/4);
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/tests/Dialect/mhlo/
Dops.mlir239 // expected-error@+1 {{AllToAll split_dimension -1 is out-of-bounds for input rank 2}}
252 // expected-error@+1 {{AllToAll split_dimension 2 is out-of-bounds for input rank 2}}
265 // expected-error@+1 {{AllToAll concat_dimension -1 is out-of-bounds for input rank 2}}
278 // expected-error@+1 {{AllToAll concat_dimension 2 is out-of-bounds for input rank 2}}
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/stablehlo/tests/
Dops_stablehlo.mlir239 // expected-error@+1 {{AllToAll split_dimension -1 is out-of-bounds for input rank 2}}
252 // expected-error@+1 {{AllToAll split_dimension 2 is out-of-bounds for input rank 2}}
265 // expected-error@+1 {{AllToAll concat_dimension -1 is out-of-bounds for input rank 2}}
278 // expected-error@+1 {{AllToAll concat_dimension 2 is out-of-bounds for input rank 2}}

12