| /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 | 74 def __init__(self, input_dataset, num_workers, index, num_replicas=None): argument
93 def _AutoShardDatasetV1(input_dataset, num_workers, index, num_replicas=None): # pylint: disable=i… argument
249 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
81 int64_t num_replicas) { in CreateConfig()
|
| 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
|
| D | auto_shard_dataset_op_test.cc | 29 int64_t num_replicas, DataTypeVector output_dtypes, in AutoShardDatasetParams()
|
| /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 | 189 int64_t num_replicas = replicate.n(); in AnnotateCompileOpAndGetExecuteArgToWhileArgsMapping() local
284 int64_t num_replicas = replicate.n(); in AddInputsToReplicateOp() local
393 int64_t num_replicas = replicate.n(); in HandleReplicateOp() local
|
| D | replicate_to_island.cc | 147 int num_replicas, llvm::SmallVectorImpl<tf_executor::IslandOp>& replicas) { in ExpandReplicateIntoReplicas()
245 const int num_replicas = replicate_op.n(); in CreateIslandsFromReplicate() local
|
| D | tpu_rewrite_pass.cc | 285 tf_device::ClusterFuncOp op, int num_replicas, int num_cores_per_replica, in SetMetadataProtoFromClusterFuncOp()
330 tf_device::ClusterFuncOp cluster_func, int num_replicas, in BuildCompileOp()
403 const int num_replicas = tpu_devices.size(); in AssignDevicesToReplicate() local
558 int num_replicas = 1; in Rewrite() local
|
| /external/tensorflow/tensorflow/compiler/xla/tools/ |
| D | hlo_module_loader.h | 30 int64 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 | 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
|
| D | tpu_compile_ops.cc | 92 tensorflow::tpu::TpuMeshStateInterface* mesh_state, int num_replicas, in CompilationCacheKey()
177 int num_replicas = config.num_replicas() ? config.num_replicas() : 1; in Compute() local
|
| /external/tensorflow/tensorflow/python/eager/benchmarks/resnet50/ |
| D | resnet50_test_util.py | 51 num_replicas=1): argument
|
| D | resnet50_test.py | 37 def compute_gradients(model, images, labels, num_replicas=1): argument
268 num_replicas=1): argument
|
| /external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/ |
| D | tpu_client_extension.cc | 64 [](PyTpuClient* client, int num_replicas) in PYBIND11_MODULE()
65 -> StatusOr<std::vector<std::shared_ptr<PjRtDevice>>> { in PYBIND11_MODULE()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_runner_interface.h | 47 int64 num_replicas = 1; member
|
| D | ar_crs_combiner.h | 75 ArCrsCombiner(int num_spatial_partitions, int num_replicas, in ArCrsCombiner()
|
| /external/tensorflow/tensorflow/core/grappler/optimizers/data/ |
| D | auto_shard.cc | 574 Status RewriteRebatchV2ToV1(const NodeDef& sink_node, int64_t num_replicas, in RewriteRebatchV2ToV1()
624 int64_t num_replicas, MutableGraphView* graph) { in ShardByData()
646 int64_t num_replicas, MutableGraphView* graph) { in ShardByHint()
677 int64_t num_replicas, GraphDef* output, in OptimizeGraph()
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | hlo_test_base.cc | 207 int64_t num_replicas, bool use_threads, bool run_hlo_passes) { in ExecuteReplicated()
220 int64_t num_replicas, DeviceAssignment* device_assignment, in ExecuteReplicated()
237 int64_t num_replicas, bool run_hlo_passes) { in ExecuteReplicated()
419 int64_t num_replicas, in RunReplicated()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | executable_build_options.cc | 63 int num_replicas) { in set_num_replicas()
|
| D | executable_build_options.h | 80 int num_replicas() const { return num_replicas_; } in num_replicas() function
|
| /external/tensorflow/tensorflow/python/distribute/parallel_device/ |
| D | parallel_device_test.py | 58 def _collective_reduce(inputs, operation, num_replicas): argument
75 def _collective_sum(inputs, num_replicas): argument
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_rewrite_device_util.cc | 188 int num_replicas, int num_cores_per_replica, in GetFullMeshTPUExecutionDeviceAssignment()
347 int num_replicas, int num_cores_per_replica, in GetGeneralTPUExecutionDeviceAssignment()
450 Devices devices, int num_replicas, int num_cores_per_replica, in GetTPUCompilationAndExecutionDevices()
|