Searched refs:replicated_inputs (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/ |
| D | tf_device.cc | 216 replicated_inputs, in ParseReplicateOpOperands() argument
245 } else if (parser->parseOperandList(replicated_inputs->emplace_back(), in ParseReplicateOpOperands()
277 replicated_inputs, in SetReplicateOpOperands() argument
288 if (replicated_inputs.empty() && packed_inputs.empty()) return success(); in SetReplicateOpOperands()
290 for (auto replicated_input_and_idx : llvm::enumerate(replicated_inputs)) { in SetReplicateOpOperands()
305 const int32_t num_replicated_block_args = replicated_inputs.size(); in SetReplicateOpOperands()
327 replicated_inputs; in ParseReplicateOp() local
333 if (ParseReplicateOpOperands(parser, state, &replicated_inputs, in ParseReplicateOp()
337 SetReplicateOpOperands(loc, parser, state, replicated_inputs, in ParseReplicateOp()
344 int32_t num_replicated_inputs = replicated_inputs.size() * n; in ParseReplicateOp()
[all …]
|
| D | tf_device_ops.td | 203 replicated_inputs: each group of `n` inputs corresponds to an input for a single
268 Variadic<AnyType>:$replicated_inputs,
299 "llvm::ArrayRef<std::pair<ValueRange, Type>>":$replicated_inputs,
302 "llvm::ArrayRef<std::pair<ValueRange, Type>>":$replicated_inputs,
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_variable_runtime_reformatting.cc | 288 llvm::SmallVector<llvm::SmallVector<Value, 8>, 8> replicated_inputs; in AddInputsToReplicateOp() local
289 replicated_inputs.reserve(replicate.GetNumReplicatedBlockArguments()); in AddInputsToReplicateOp()
292 replicated_inputs.emplace_back(); in AddInputsToReplicateOp()
294 replicated_inputs.back().push_back( in AddInputsToReplicateOp()
297 new_replicated_inputs.emplace_back(replicated_inputs.back(), arg.getType()); in AddInputsToReplicateOp()
|
| D | tpu_cluster_formation.cc | 393 llvm::SmallVector<std::pair<ValueRange, Type>, 8> replicated_inputs; in ReplicateCluster() local
409 replicated_inputs.push_back( in ReplicateCluster()
427 replicated_inputs, packed_inputs, cluster.getResultTypes()); in ReplicateCluster()
|
| /external/tensorflow/tensorflow/python/keras/distribute/ |
| D | custom_training_loop_models_test.py | 101 def train_step(replicated_inputs): argument
111 outputs = distribution.run(step_fn, args=(replicated_inputs,))
|