TPUReplicateMetadata.Options

public static class TPUReplicateMetadata.Options

Optional attributes for TPUReplicateMetadata

Public Methods

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

Inherited Methods

Public Methods

public TPUReplicateMetadata.Options allowSoftPlacement (Boolean allowSoftPlacement)

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

Parameters
computationShape DEPRECATED. Use num_cores_per_replica instead.

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

Parameters
deviceAssignment The assignment of devices for the computation.

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

public TPUReplicateMetadata.Options numCoresPerReplica (Long numCoresPerReplica)

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

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

public TPUReplicateMetadata.Options stepMarkerLocation (String stepMarkerLocation)

public TPUReplicateMetadata.Options topology (String topology)

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

public TPUReplicateMetadata.Options useSpmdForXlaPartitioning (Boolean useSpmdForXlaPartitioning)

public TPUReplicateMetadata.Options useTpu (Boolean useTpu)

Parameters
useTpu Whether to place the computation on the TPU.