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. |
© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.4/api_docs/python/tf/raw_ops/TPUReplicateMetadata