Home
last modified time | relevance | path

Searched defs:num_replicas (Results 1 – 25 of 77) sorted by relevance

1234

/aosp_15_r20/external/tensorflow/tensorflow/python/data/experimental/ops/
H A Ddistribute.py70 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
/aosp_15_r20/external/tensorflow/tensorflow/core/grappler/optimizers/data/
H A Dauto_shard.cc655 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()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dgpu_spmd_partitioner.h31 HloComputation* computation, int64_t num_partitions, int64_t num_replicas, in GpuSpmdPartitioningVisitor()
44 GpuSpmdPartitioner(int64_t num_partitions, int64_t num_replicas) in GpuSpmdPartitioner()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dstateful_rng_spmd_partitioner.h32 HloComputation* computation, int64_t num_partitions, int64_t num_replicas, in StatefulRngSpmdPartitioningVisitor()
45 StatefulRngSpmdPartitioner(int64_t num_partitions, int64_t num_replicas) in StatefulRngSpmdPartitioner()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
H A Dtpu_variable_runtime_reformatting.cc133 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
H A Dtpu_rewrite_pass.cc280 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
H A Dreplicate_invariant_op_hoisting.cc46 void MakeShapeOpInvariant(tf_device::ReplicateOp replicate_op, int num_replicas, in MakeShapeOpInvariant()
125 const int num_replicas = replicate_op.n(); in HoistReplicateInvariantOps() local
H A Dreplicate_to_island.cc143 int num_replicas, llvm::SmallVectorImpl<tf_executor::IslandOp>& replicas) { in ExpandReplicateIntoReplicas()
241 const int num_replicas = replicate_op.n(); in CreateIslandsFromReplicate() local
H A Dtpu_cluster_formation.cc381 LogicalResult ReplicateCluster(tf_device::ClusterOp cluster, int num_replicas, in ReplicateCluster()
628 int num_replicas = num_replicas_attr.cast<mlir::IntegerAttr>().getInt(); in FormClustersInBlock() local
/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/data/experimental/
H A Dauto_shard_dataset_op.cc47 int64_t index, num_workers, auto_shard_policy, num_replicas; in MakeDataset() local
83 int64_t num_replicas) { in CreateConfig()
H A Drebatch_dataset_op.cc44 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
H A Dauto_shard_dataset_op_test.cc34 int64_t num_replicas, DataTypeVector output_dtypes, in AutoShardDatasetParams()
288 Node* num_replicas; in TEST_F() local
/aosp_15_r20/external/pytorch/torch/csrc/api/src/data/samplers/
H A Ddistributed.cpp17 size_t num_replicas, in DistributedRandomSampler()
99 size_t num_replicas, in DistributedSequentialSampler()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/
H A Dhlo_test_base.cc250 int64_t num_replicas, bool use_threads, bool run_hlo_passes) { in ExecuteReplicated()
263 int64_t num_replicas, DeviceAssignment* device_assignment, in ExecuteReplicated()
280 int64_t num_replicas, bool run_hlo_passes, in ExecuteReplicated()
566 string_view hlo_string, bool run_hlo_passes, int64_t num_replicas, in RunReplicated()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/kernels/
H A Dxrt_compile_ops.cc55 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
H A Dtpu_compile_ops.cc93 tensorflow::tpu::TpuMeshStateInterface* mesh_state, int num_replicas, in CompilationCacheKey()
178 int num_replicas = config.num_replicas() ? config.num_replicas() : 1; in Compute() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
H A Dtpu_rewrite_device_util.cc184 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()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
H A Dtpu_client_extension.cc69 [](PyTpuClient* client, int num_replicas) in PYBIND11_MODULE()
70 -> StatusOr<std::vector<std::shared_ptr<PjRtDevice>>> { in PYBIND11_MODULE()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/graph_rewrite/
H A Dencapsulate_tpu_computations_pass.cc431 int num_replicas; in MoveHeadOutsideCompilationToHost() local
750 int num_replicas; in RemoveUnusedXlaInput() local
1023 int num_replicas; in MoveTailOutsideCompilationToHost() local
1211 int num_replicas; in ReplaceArgUsedByOutsideCompilationWithPlaceholder() local
1333 int num_replicas; in RemoveUnusedXlaOutput() local
1441 int num_replicas; in RemoveEdgesBetweenArgAndRetval() local
2553 int num_replicas; in BuildTPUReplicateOps() local
H A Ddistributed_tpu_rewrite_pass.cc1034 const xla::OpSharding& sharding, const int num_replicas, in CreateOrGetXlaSplitNodeForDistributedArg()
1064 const xla::OpSharding& sharding, const int num_replicas, in CreateOrGetXlaSplitNodeForVariableArg()
1653 const tpu::TpuTopologyExternal& tpu_topology, int num_replicas, in ParseDeviceAssignmentAttr()
1713 int num_replicas, const std::vector<std::vector<Device*>>& tpu_devices, in BuildFullMeshDeviceAssignment()
1736 int num_replicas, int num_cores_per_replica, in BuildGeneralDeviceAssignment()
1777 const std::vector<std::vector<Device*>>& tpu_devices, int num_replicas, in BuildDeviceAssignment()
3076 const xla::OpSharding& sharding, const int num_replicas, in CreatePartitionedDummyVarArgs()
4372 const DeviceSet& device_set, const Node& replicate_node, int* num_replicas, in GetDeviceTopology()
4427 int num_replicas, const Node& replicate_node, FunctionLibraryRuntime* flr, in GetIOTypes()
4725 int num_replicas; in RewriteTPUReplicateNode() local
/aosp_15_r20/external/tensorflow/tensorflow/python/eager/benchmarks/resnet50/
H A Dresnet50_test.py32 def compute_gradients(model, images, labels, num_replicas=1): argument
263 num_replicas=1): argument
/aosp_15_r20/external/tensorflow/tensorflow/python/distribute/parallel_device/
H A Dparallel_device_test.py54 def _collective_reduce(inputs, operation, num_replicas): argument
71 def _collective_sum(inputs, num_replicas): argument
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/
H A Dpjrt_c_api_client.h146 int num_replicas, int num_partitions) const override { in GetDefaultDeviceAssignment()
433 int num_replicas() const override { return wrapped()->num_replicas(); } in num_replicas() function
H A Dtfrt_cpu_pjrt_client.cc233 int num_replicas, int num_partitions) const { in GetDefaultDeviceAssignment()
337 int num_replicas; in Compile() local
342 [this](int num_replicas, int num_partitions) { in Compile()
1158 int num_replicas, int num_partitions, in TfrtCpuExecutable()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dreduce_scatter_utils.cc254 int64_t num_replicas, bool allow_multiple_split_dims, in MatchReduceScatter()
270 int64_t num_replicas, bool allow_multiple_split_dims, in MatchReduceScatter()

1234