| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_test.cc | 661 auto left_halo = AllOf(op::CollectivePermute(op::Slice(lhs)), in TEST_F()
663 auto right_halo = AllOf(op::CollectivePermute(op::Slice(lhs)), in TEST_F()
900 op::CollectivePermute(op::Slice(sharded_input))); in TEST_F()
936 op::CollectivePermute(op::Slice(op::Parameter(0)))); in TEST_F()
938 AllOf(op::Shape("f32[2,2]"), op::CollectivePermute(op::Parameter(0))); in TEST_F()
978 op::CollectivePermute(op::Slice(sharded_input))); in TEST_F()
1020 op::CollectivePermute(op::Slice(sharded_input))); in TEST_F()
1022 op::CollectivePermute(op::Slice(sharded_input))); in TEST_F()
1069 op::CollectivePermute(op::Slice(sharded_input))); in TEST_F()
1071 op::CollectivePermute(op::Slice(sharded_input))); in TEST_F()
[all …]
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v2/ |
| D | CollectivePermute.pbtxt | 2 name: "CollectivePermute"
42 name: "CollectivePermute"
|
| /external/tensorflow/tensorflow/core/ops/compat/ops_history_v1/ |
| D | CollectivePermute.pbtxt | 2 name: "CollectivePermute"
|
| /external/tensorflow/tensorflow/core/api_def/base_api/ |
| D | api_def_CollectivePermute.pbtxt | 2 graph_op_name: "CollectivePermute"
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | cross_replica_ops.cc | 120 xla::CollectivePermute(ctx->Input(0), source_target_pairs)); in Compile()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | jitrt_custom_calls.cc | 1988 struct CollectivePermute { struct
1998 static CollectivePermute Handler() { return CollectivePermute(); } in Handler() argument
2002 Error CollectivePermute::operator()( in operator ()()
2065 static bool CollectivePermute(runtime::KernelContext* ctx, void** args, in CollectivePermute() function
2079 .To<RuntimeChecks()>(CollectivePermute::Handler()) in CollectivePermute()
2175 lib.Insert("xla.gpu.collective_permute", &xla::gpu::CollectivePermute); in JitRtGpuCustomCalls()
|
| /external/tensorflow/tensorflow/compiler/xla/mlir_hlo/include/mlir-hlo/Dialect/mhlo/IR/ |
| D | hlo_ops_base_attrs.td | 112 // optionally for collective instructions (AllReduce, CollectivePermute,
|
| D | hlo_ops.td | 1645 let summary = "CollectivePermute operator";
1647 CollectivePermute is a collective operation that sends and receives data
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | simple_orc_jit.cc | 277 REGISTER_CPU_RUNTIME_SYMBOL(CollectivePermute); in RegisterKnownJITSymbols()
|
| /external/tensorflow/tensorflow/compiler/xla/mlir_hlo/stablehlo/dialect/ |
| D | StablehloAttrs.td | 116 // optionally for collective instructions (AllReduce, CollectivePermute,
|
| D | StablehloOps.td | 1521 let summary = "CollectivePermute operator";
1523 CollectivePermute is a collective operation that sends and receives data
|
| /external/tensorflow/tensorflow/compiler/xla/tools/ |
| D | hlo_control_flow_flattening_test.cc | 418 TEST_F(HloControlFlowFlatteningTest, CollectivePermute) { in TEST_F() argument
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_matchers.h | 254 HLO_MATCHER(CollectivePermute);
|
| D | hlo_parser_test.cc | 1862 "CollectivePermute", in CreateTestCases()
1863 R"(HloModule CollectivePermute, entry_computation_layout={(f32[128,32]{0,1})->f32[128,32]{0,1}} in CreateTestCases()
1865 ENTRY CollectivePermute { in CreateTestCases()
|
| D | hlo.proto | 185 // optionally for collective instructions (AllReduce, CollectivePermute,
|
| /external/tensorflow/tensorflow/compiler/xla/python/xla_extension/ |
| D | ops.pyi | 133 def CollectivePermute(
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | ops.cc | 123 ops.def("CollectivePermute", &CollectivePermute, py::arg("operand"), in BuildOpsSubmodule()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder_test.cc | 489 TEST_F(XlaBuilderTest, CollectivePermute) { in TEST_F() argument
492 CollectivePermute(x, {{0, 1}, {1, 2}, {2, 3}}); in TEST_F()
|
| D | xla_builder.h | 783 XlaOp CollectivePermute(
1392 friend XlaOp CollectivePermute(
2390 XlaOp CollectivePermute(
|
| D | xla_builder.cc | 3226 XlaOp XlaBuilder::CollectivePermute( in CollectivePermute() function in xla::XlaBuilder
4714 XlaOp CollectivePermute( in CollectivePermute() function
4717 return operand.builder()->CollectivePermute(operand, source_target_pairs); in CollectivePermute()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf_patterns.td | 295 // CollectivePermute op patterns.
|
| /external/tensorflow/tensorflow/compiler/xla/mlir_hlo/include/mlir-hlo/Dialect/lhlo/IR/ |
| D | lhlo_ops.td | 1147 let summary = "CollectivePermute operator";
1149 CollectivePermute is a collective operation that sends and receives data
|
| /external/tensorflow/tensorflow/compiler/xla/g3doc/ |
| D | operation_semantics.md | 603 ## CollectivePermute section in Operation Semantics
606 [`XlaBuilder::CollectivePermute`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xl…
608 CollectivePermute is a collective operation that sends and receives data cross
611 <b> `CollectivePermute(operand, source_target_pairs)` </b>
|
| /external/tensorflow/tensorflow/tools/api/golden/v2/ |
| D | tensorflow.raw_ops.pbtxt | 828 name: "CollectivePermute"
|
| /external/tensorflow/tensorflow/tools/api/golden/v1/ |
| D | tensorflow.raw_ops.pbtxt | 828 name: "CollectivePermute"
|