View source on GitHub |
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.
Returns | |
---|---|
A mesh created from specified or default arguments. |