tf.experimental.dtensor.create_distributed_mesh

Creates a single- or multi-client mesh.

tf.experimental.dtensor.create_distributed_mesh(
    mesh_dims: List[Tuple[str, int]],
    mesh_name: str = '',
    num_global_devices: Optional[int] = None,
    num_clients: Optional[int] = None,
    client_id: Optional[int] = None,
    device_type: str = 'CPU'
) -> tf.experimental.dtensor.Mesh

For CPU and GPU meshes, users can choose to use fewer local devices than what is available. If any argument is missing, it will be extracted from environment variables. The default values for these environment variables create a mesh using all devices (common for unit tests).

For TPU meshes, users should not specify any of the nullable arguments. The DTensor runtime will set these arguments automatically, using all TPU cores available in the entire cluster.

mesh_dims A list of (dim_name, dim_size) tuples.
mesh_name Name of the created mesh. Defaults to ''.
num_global_devices Number of devices in the DTensor cluster. Defaults to the corresponding environment variable.
num_clients Number of clients in the DTensor cluster. Defaults to the corresponding environment variable, DTENSOR_NUM_CLIENTS.
client_id This client's ID. Defaults to the corresponding environment variable, DTENSOR_CLIENT_ID.
device_type Type of device to build the mesh for. Defaults to 'CPU'.

A mesh created from specified or default arguments.