public final class
TPUReplicateMetadata
Metadata indicating how the TPU computation should be replicated.
This operation holds the metadata common to operations of a `tpu.replicate()` computation subgraph.
Nested Classes
class | TPUReplicateMetadata.Options |
Optional attributes for
TPUReplicateMetadata
|
Public Methods
static TPUReplicateMetadata.Options |
allowSoftPlacement
(Boolean allowSoftPlacement)
|
static TPUReplicateMetadata.Options |
computationShape
(List<Long> computationShape)
|
static TPUReplicateMetadata |
create
(
Scope
scope, Long numReplicas,
Options...
options)
Factory method to create a class wrapping a new TPUReplicateMetadata operation.
|
static TPUReplicateMetadata.Options |
deviceAssignment
(List<Long> deviceAssignment)
|
static TPUReplicateMetadata.Options |
hostComputeCore
(List<String> hostComputeCore)
|
static TPUReplicateMetadata.Options |
numCoresPerReplica
(Long numCoresPerReplica)
|
static TPUReplicateMetadata.Options |
paddingMap
(List<String> paddingMap)
|
static TPUReplicateMetadata.Options |
stepMarkerLocation
(String stepMarkerLocation)
|
static TPUReplicateMetadata.Options |
topology
(String topology)
|
static TPUReplicateMetadata.Options |
useSpmdForXlaPartitioning
(Boolean useSpmdForXlaPartitioning)
|
static TPUReplicateMetadata.Options |
useTpu
(Boolean useTpu)
|
Inherited Methods
Public Methods
public static TPUReplicateMetadata.Options computationShape (List<Long> computationShape)
Parameters
computationShape | DEPRECATED. Use num_cores_per_replica instead. |
---|
public static TPUReplicateMetadata create ( Scope scope, Long numReplicas, Options... options)
Factory method to create a class wrapping a new TPUReplicateMetadata operation.
Parameters
scope | current scope |
---|---|
numReplicas | Number of replicas of the computation |
options | carries optional attributes values |
Returns
- a new instance of TPUReplicateMetadata
public static TPUReplicateMetadata.Options deviceAssignment (List<Long> deviceAssignment)
Parameters
deviceAssignment | The assignment of devices for the computation. |
---|
public static TPUReplicateMetadata.Options numCoresPerReplica (Long numCoresPerReplica)
Parameters
numCoresPerReplica | Number of cores per replica. Used for model parallelism. |
---|
public static TPUReplicateMetadata.Options topology (String topology)
Parameters
topology | TopologyProto indicating the topology of the TPU pod slice. |
---|
public static TPUReplicateMetadata.Options useSpmdForXlaPartitioning (Boolean useSpmdForXlaPartitioning)
public static TPUReplicateMetadata.Options useTpu (Boolean useTpu)
Parameters
useTpu | Whether to place the computation on the TPU. |
---|