Module: tf.experimental.dtensor

Public API for tf.experimental.dtensor namespace.

Classes

class DTensorCheckpoint: Manages saving/restoring trackable values to disk, for DTensor.

class DVariable: A replacement for tf.Variable which follows initial value placement.

class Layout: Represents the layout information of a DTensor.

class Mesh: Represents a Mesh configuration over a certain list of Mesh Dimensions.

Functions

call_with_layout(...): Calls a function in the DTensor device scope if layout is not None.

check_layout(...): Asserts that the layout of the DTensor is layout.

client_id(...): Returns this client's ID.

copy_to_mesh(...): Copies a tf.Tensor onto the DTensor device with the given layout.

create_distributed_mesh(...): Creates a single- or multi-client mesh.

create_mesh(...): Creates a single-client mesh.

device_name(...): Returns the singleton DTensor device's name.

enable_save_as_bf16(...): Allows float32 DVariables to be checkpointed and restored as bfloat16.

fetch_layout(...): Fetches the layout of a DTensor.

full_job_name(...): Returns the fully qualified TF job name for this or another task.

heartbeat_enabled(...): Returns true if DTensor heartbeat service is enabled.

initialize_multi_client(...): Initializes Multi Client DTensor.

initialize_tpu_system(...): Initialize the TPU devices.

job_name(...): Returns the job name used by all clients in this DTensor cluster.

jobs(...): Returns a list of job names of all clients in this DTensor cluster.

local_devices(...): Returns a list of device specs of device_type attached to this client.

name_based_restore(...): Restores from checkpoint_prefix to name based DTensors.

name_based_save(...): Saves name based Tensor into a Checkpoint.

num_clients(...): Returns the number of clients in this DTensor cluster.

num_global_devices(...): Returns the number of devices of device_type in this DTensor cluster.

num_local_devices(...): Returns the number of devices of device_type attached to this client.

pack(...): Packs tf.Tensor components into a DTensor.

relayout(...): Changes the layout of tensor.

run_on(...): Runs enclosed functions in the DTensor device scope.

sharded_save(...): Saves given named tensor slices in a sharded, multi-client safe fashion.

shutdown_tpu_system(...): Shutdown TPU system.

unpack(...): Unpacks a DTensor into tf.Tensor components.

MATCH 'match'
UNSHARDED 'unsharded'