tensorflow / 2.9.1 / experimental / dtensor / create_mesh.html /

tf.experimental.dtensor.create_mesh

Creates a single-client 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