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 guide | Used in the tutorials |
|---|---|
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. |
© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/experimental/dtensor/create_mesh