Home
last modified time | relevance | path

Searched refs:replicated_input (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dtpu_reorder_replicate_and_partitioned_inputs.cc34 TF::TPUReplicatedInputOp replicated_input) { in ReorderReplicateAndPartitionedInputs() argument
35 if (!llvm::all_of(replicated_input.inputs(), [](Value input) { in ReorderReplicateAndPartitionedInputs()
39 return replicated_input.emitOpError() in ReorderReplicateAndPartitionedInputs()
43 replicated_input.getOperand(0).getDefiningOp()); in ReorderReplicateAndPartitionedInputs()
49 for (auto operand : replicated_input.inputs().drop_front()) { in ReorderReplicateAndPartitionedInputs()
65 return replicated_input.emitOpError() in ReorderReplicateAndPartitionedInputs()
78 for (auto operand : replicated_input.inputs()) { in ReorderReplicateAndPartitionedInputs()
88 OpBuilder builder(replicated_input); in ReorderReplicateAndPartitionedInputs()
92 replicated_input.getLoc(), replicated_input.getType(), in ReorderReplicateAndPartitionedInputs()
93 operands_per_replica, replicated_input->getAttrs()); in ReorderReplicateAndPartitionedInputs()
[all …]
Dtf_passes.td832 %replicated_input = "tf.TPUReplicatedInput"(%arg0) : (tensor<i32>) -> tensor<i32>
833 …%identity = "tf.Identity"(%replicated_input) {_xla_compile_device_type = "TPU", _replication_info …
856 …%replicated_input = "tf.TPUReplicatedInput"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i…
857 …%identity = "tf.Identity"(%replicated_input) {_xla_compile_device_type = "TPU", _replication_info …
867 %replicate:2 = tf_device.replicate([%arg0, %arg1] as %replicated_input) {n = 2 : i32} {
869 %identity = "tf.Identity"(%replicated_input) : (tensor<i32>) -> tensor<i32>
/external/tensorflow/tensorflow/core/api_def/base_api/
Dapi_def_TPUReplicatedInput.pbtxt13 %replicated_input = "tf.TPUReplicatedInput"(%a, %b)
14 %computation = "tf.Computation"(%replicated_input)
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/
Dtf_device.cc294 const auto& replicated_input = replicated_input_and_idx.value(); in SetReplicateOpOperands() local
296 if (replicated_input.size() != *n) in SetReplicateOpOperands()
299 << " to be 'n' (" << *n << "), got " << replicated_input.size(); in SetReplicateOpOperands()
302 if (parser->resolveOperands(replicated_input, region_arg_types[idx], in SetReplicateOpOperands()
448 for (auto& replicated_input : replicated_inputs) { in BuildReplicateOp() local
449 DCHECK_EQ(llvm::size(replicated_input.first), n); in BuildReplicateOp()
450 for (auto input : replicated_input.first) { in BuildReplicateOp()
452 VerifyCompatibleTypes(input.getType(), replicated_input.second))); in BuildReplicateOp()
455 block.addArgument(replicated_input.second, state->location); in BuildReplicateOp()
Dtf_generated_ops.td17770 %replicated_input = "tf.TPUReplicatedInput"(%a, %b)
17771 %computation = "tf.Computation"(%replicated_input)
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_functional_ops.cc842 std::vector<NodeBuilder::NodeOut> replicated_input; in CreateConcatAndSplitNodesForInputTensor() local
843 replicated_input.push_back(NodeBuilder::NodeOut(concat_node)); in CreateConcatAndSplitNodesForInputTensor()
849 .Input(replicated_input) in CreateConcatAndSplitNodesForInputTensor()