ReplicateMetadata.Options

public static class ReplicateMetadata.Options

Optional attributes for ReplicateMetadata

Public Methods

ReplicateMetadata.Options
allowSoftPlacement(Boolean allowSoftPlacement)
ReplicateMetadata.Options
computationShape(List<Long> computationShape)
ReplicateMetadata.Options
deviceAssignment(List<Long> deviceAssignment)
ReplicateMetadata.Options
hostComputeCore(List<String> hostComputeCore)
ReplicateMetadata.Options
numCoresPerReplica(Long numCoresPerReplica)
ReplicateMetadata.Options
paddingMap(List<String> paddingMap)
ReplicateMetadata.Options
stepMarkerLocation(String stepMarkerLocation)
ReplicateMetadata.Options
topology(String topology)
ReplicateMetadata.Options
useSpmdForXlaPartitioning(Boolean useSpmdForXlaPartitioning)
ReplicateMetadata.Options
useTpu(Boolean useTpu)

Inherited Methods

Public Methods

public ReplicateMetadata.Options allowSoftPlacement (Boolean allowSoftPlacement)

public ReplicateMetadata.Options computationShape (List<Long> computationShape)

Parameters
computationShape DEPRECATED. Use num_cores_per_replica instead.

public ReplicateMetadata.Options deviceAssignment (List<Long> deviceAssignment)

Parameters
deviceAssignment The assignment of devices for the computation.

public ReplicateMetadata.Options hostComputeCore (List<String> hostComputeCore)

public ReplicateMetadata.Options numCoresPerReplica (Long numCoresPerReplica)

Parameters
numCoresPerReplica Number of cores per replica. Used for model parallelism.

public ReplicateMetadata.Options paddingMap (List<String> paddingMap)

public ReplicateMetadata.Options stepMarkerLocation (String stepMarkerLocation)

public ReplicateMetadata.Options topology (String topology)

Parameters
topology TopologyProto indicating the topology of the TPU pod slice.

public ReplicateMetadata.Options useSpmdForXlaPartitioning (Boolean useSpmdForXlaPartitioning)

public ReplicateMetadata.Options useTpu (Boolean useTpu)

Parameters
useTpu Whether to place the computation on the TPU.