tf.raw_ops.TPUReplicateMetadata
Metadata indicating how the TPU computation should be replicated.
tf.raw_ops.TPUReplicateMetadata(
num_replicas,
num_cores_per_replica=1,
topology='',
use_tpu=True,
device_assignment=[],
computation_shape=[],
host_compute_core=[],
padding_map=[],
step_marker_location='STEP_MARK_AT_ENTRY',
allow_soft_placement=False,
use_spmd_for_xla_partitioning=False,
name=None
)
This operation holds the metadata common to operations of a tpu.replicate() computation subgraph.
| Args | |
|---|---|
num_replicas |
An int that is >= 0. Number of replicas of the computation |
num_cores_per_replica |
An optional int. Defaults to 1. Number of cores per replica. Used for model parallelism. |
topology |
An optional string. Defaults to "". TopologyProto indicating the topology of the TPU pod slice. |
use_tpu |
An optional bool. Defaults to True. Whether to place the computation on the TPU. |
device_assignment |
An optional list of ints. Defaults to []. The assignment of devices for the computation. |
computation_shape |
An optional list of ints. Defaults to []. DEPRECATED. Use num_cores_per_replica instead. |
host_compute_core |
An optional list of strings. Defaults to []. |
padding_map |
An optional list of strings. Defaults to []. |
step_marker_location |
An optional string. Defaults to "STEP_MARK_AT_ENTRY". |
allow_soft_placement |
An optional bool. Defaults to False. |
use_spmd_for_xla_partitioning |
An optional bool. Defaults to False. |
name |
A name for the operation (optional). |
| Returns | |
|---|---|
| The created Operation. |
© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/raw_ops/TPUReplicateMetadata