tf.experimental.dtensor.create_mesh
Creates a single-client mesh.
tf.experimental.dtensor.create_mesh( mesh_dims: Optional[List[Tuple[str, int]]] = None, mesh_name: str = '', devices: Optional[List[str]] = None, device_type: Optional[str] = None ) -> tf.experimental.dtensor.Mesh
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 list of (dim_name, dim_size) tuples. Defaults to a single batch-parallel dimension called 'x' using all devices. As a special case, a single-element mesh_dims whose dim_size is -1 also uses all devices. |
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'. |
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/versions/r2.9/api_docs/python/tf/experimental/dtensor/create_mesh