public final class TPUReplicateMetadata extends PrimitiveOp
Modifier and Type | Class and Description |
---|---|
static class |
TPUReplicateMetadata.Options
Optional attributes for
TPUReplicateMetadata |
operation
Modifier and Type | Method and Description |
---|---|
static TPUReplicateMetadata.Options |
allowSoftPlacement(Boolean allowSoftPlacement) |
static TPUReplicateMetadata.Options |
computationShape(List<Long> computationShape) |
static TPUReplicateMetadata |
create(Scope scope,
Long numReplicas,
TPUReplicateMetadata.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 |
useTpu(Boolean useTpu) |
equals, hashCode, op, toString
public static TPUReplicateMetadata create(Scope scope, Long numReplicas, TPUReplicateMetadata.Options... options)
scope
- current scopenumReplicas
- Number of replicas of the computationoptions
- carries optional attributes valuespublic static TPUReplicateMetadata.Options numCoresPerReplica(Long numCoresPerReplica)
numCoresPerReplica
- Number of cores per replica. Used for model parallelism.public static TPUReplicateMetadata.Options topology(String topology)
topology
- TopologyProto indicating the topology of the TPU pod slice.public static TPUReplicateMetadata.Options useTpu(Boolean useTpu)
useTpu
- Whether to place the computation on the TPU.public static TPUReplicateMetadata.Options deviceAssignment(List<Long> deviceAssignment)
deviceAssignment
- The assignment of devices for the computation.public static TPUReplicateMetadata.Options computationShape(List<Long> computationShape)
computationShape
- DEPRECATED. Use num_cores_per_replica instead.public static TPUReplicateMetadata.Options hostComputeCore(List<String> hostComputeCore)
hostComputeCore
- public static TPUReplicateMetadata.Options paddingMap(List<String> paddingMap)
paddingMap
- public static TPUReplicateMetadata.Options stepMarkerLocation(String stepMarkerLocation)
stepMarkerLocation
- public static TPUReplicateMetadata.Options allowSoftPlacement(Boolean allowSoftPlacement)
allowSoftPlacement
- Copyright © 2022. All rights reserved.