tf.experimental.dtensor.create_mesh

Creates a single-client mesh.

tf.experimental.dtensor.create_mesh(
    mesh_dims: Optional[Union[List[Tuple[str, int]], Dict[str, int]]] = None,
    mesh_name: str = '',
    devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
    device_type: Optional[str] = None,
    use_xla_spmd: bool = layout.USE_XLA_SPMD
) -> tf.experimental.dtensor.Mesh

Used in the notebooks

Used in the guide	Used in the tutorials
DTensor concepts Distributed training with Core APIs and DTensor	Distributed training with DTensors Using DTensors with Keras

If both mesh_dims and devices are specified, they must match each otehr. As a special case, when all arguments are missing, this creates a 1D CPU mesh with an empty name, assigning all available devices to that dimension.

Args
mesh_dims	A dict of dim_name: dim_size, or a list of (dim_name, dim_size) tuples. Defaults to a single batch-parallel dimension called 'x' usin all devices. As a special case, a single-element mesh_dims whose dim_size is -1 also uses all devices. e.g. {'x' : 4, 'y' : 1} or [('x', 4), ('y', 1)].
mesh_name	Name of the created mesh. Defaults to ''.
devices	String representations of devices to use. This is the device part of tf.DeviceSpec, e.g. 'CPU:0'. Defaults to all available logical devices.
device_type	If devices is missing, the type of devices to use. Defaults to 'CPU'.
use_xla_spmd	Boolean when True, will use XLA SPMD instead of DTensor SPMD.

Returns
A single-client mesh created from specified or default arguments.