| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | stateful_rng_spmd_partitioner.h | 32 HloComputation* computation, int64_t num_partitions, int64_t num_replicas, in StatefulRngSpmdPartitioningVisitor()
45 StatefulRngSpmdPartitioner(int64_t num_partitions, int64_t num_replicas) in StatefulRngSpmdPartitioner()
|
| D | stateful_rng_spmd_partitioner.cc | 50 HloComputation* computation, int64_t num_partitions, int64_t num_replicas, in CreateVisitor()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | gpu_spmd_partitioner.h | 31 HloComputation* computation, int64_t num_partitions, int64_t num_replicas, in GpuSpmdPartitioningVisitor()
44 GpuSpmdPartitioner(int64_t num_partitions, int64_t num_replicas) in GpuSpmdPartitioner()
|
| D | gpu_spmd_partitioner.cc | 50 HloComputation* computation, int64_t num_partitions, int64_t num_replicas, in CreateVisitor()
|
| /external/tensorflow/tensorflow/python/data/experimental/ops/ |
| D | distribute.py | 70 def __init__(self, input_dataset, num_workers, index, num_replicas=None): argument
89 def _AutoShardDatasetV1(input_dataset, num_workers, index, num_replicas=None): # pylint: disable=i… argument
258 def __init__(self, input_dataset, num_replicas): argument
|
| /external/tensorflow/tensorflow/core/kernels/data/experimental/ |
| D | auto_shard_dataset_op.cc | 47 int64_t index, num_workers, auto_shard_policy, num_replicas; in MakeDataset() local
83 int64_t num_replicas) { in CreateConfig()
|
| D | auto_shard_dataset_op_test.cc | 34 int64_t num_replicas, DataTypeVector output_dtypes, in AutoShardDatasetParams()
288 Node* num_replicas; in TEST_F() local
|
| D | rebatch_dataset_op.cc | 44 int64_t num_replicas; in MakeDataset() local
58 const int64_t num_replicas, const DataTypeVector& output_types, in Dataset()
110 Node* num_replicas = nullptr; in AsGraphDefInternal() local
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | replicate_invariant_op_hoisting.cc | 46 void MakeShapeOpInvariant(tf_device::ReplicateOp replicate_op, int num_replicas, in MakeShapeOpInvariant()
125 const int num_replicas = replicate_op.n(); in HoistReplicateInvariantOps() local
|
| D | tpu_variable_runtime_reformatting.cc | 133 int64_t num_replicas = replicate.n(); in AnnotateCompileOpAndGetExecuteArgToWhileArgsMapping() local
228 int64_t num_replicas = replicate.n(); in AddInputsToReplicateOp() local
337 int64_t num_replicas = replicate.n(); in HandleReplicateOp() local
|
| D | replicate_to_island.cc | 143 int num_replicas, llvm::SmallVectorImpl<tf_executor::IslandOp>& replicas) { in ExpandReplicateIntoReplicas()
241 const int num_replicas = replicate_op.n(); in CreateIslandsFromReplicate() local
|
| D | tpu_rewrite_pass.cc | 280 tf_device::ClusterFuncOp op, int num_replicas, int num_cores_per_replica, in SetMetadataProtoFromClusterFuncOp()
332 tf_device::ClusterFuncOp cluster_func, int num_replicas, in BuildCompileOp()
408 const int num_replicas = tpu_devices.size(); in AssignDevicesToReplicate() local
719 int num_replicas = 1; in Rewrite() local
|
| /external/tensorflow/tensorflow/compiler/xla/tools/ |
| D | hlo_module_loader.h | 31 int64_t num_replicas = 1; member
|
| /external/tensorflow/tensorflow/core/grappler/optimizers/ |
| D | auto_parallel.h | 29 AutoParallel(int num_replicas) : num_replicas_(num_replicas) { in AutoParallel()
|
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | tpu_compile_ops.cc | 93 tensorflow::tpu::TpuMeshStateInterface* mesh_state, int num_replicas, in CompilationCacheKey()
178 int num_replicas = config.num_replicas() ? config.num_replicas() : 1; in Compute() local
|
| D | xrt_compile_ops.cc | 55 const xrt::DeviceAssignment& xrt_device_assignment, int num_replicas, in GenerateXlaDeviceAssignment()
133 int num_replicas = config.num_replicas() ? config.num_replicas() : 1; in Compile() local
|
| /external/tensorflow/tensorflow/python/eager/benchmarks/resnet50/ |
| D | resnet50_test_util.py | 47 num_replicas=1): argument
|
| D | resnet50_test.py | 32 def compute_gradients(model, images, labels, num_replicas=1): argument
263 num_replicas=1): argument
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | compile_options.proto | 31 int64 num_replicas = 4; field
|
| D | pjrt_c_api_client.h | 146 int num_replicas, int num_partitions) const override { in GetDefaultDeviceAssignment()
433 int num_replicas() const override { return wrapped()->num_replicas(); } in num_replicas() function
|
| /external/tensorflow/tensorflow/core/grappler/optimizers/data/ |
| D | auto_shard.cc | 655 Status RewriteRebatchV2ToV1(const NodeDef& sink_node, int64_t num_replicas, in RewriteRebatchV2ToV1()
705 int64_t num_replicas, MutableGraphView* graph) { in ShardByData()
727 int64_t num_replicas, MutableGraphView* graph) { in ShardByHint()
758 int64_t num_replicas, MutableGraphView* graph, in ApplyAutoShard()
790 int64_t num_replicas, GraphDef* output) { in OptimizeGraph()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_runner_interface.h | 47 int64_t num_replicas = 1; member
|
| /external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/ |
| D | tpu_client_extension.cc | 69 [](PyTpuClient* client, int num_replicas) in PYBIND11_MODULE()
70 -> StatusOr<std::vector<std::shared_ptr<PjRtDevice>>> { in PYBIND11_MODULE()
|
| /external/tensorflow/tensorflow/python/distribute/parallel_device/ |
| D | parallel_device_test.py | 54 def _collective_reduce(inputs, operation, num_replicas): argument
71 def _collective_sum(inputs, num_replicas): argument
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_rewrite_device_util.cc | 184 int num_replicas, int num_cores_per_replica, in GetFullMeshTPUExecutionDeviceAssignment()
343 int num_replicas, int num_cores_per_replica, in GetGeneralTPUExecutionDeviceAssignment()
446 Devices devices, int num_replicas, int num_cores_per_replica, in GetTPUCompilationAndExecutionDevices()
|