tf.experimental.dtensor.default_mesh

Sets the default DTensor device mesh to use for enclosed functions.

@contextlib.contextmanager
tf.experimental.dtensor.default_mesh(
    mesh: tf.experimental.dtensor.Mesh
)

This function returns a scope. All the ops and tf.functions in this scope will default to this DTensor mesh if a mesh cannot be inferred from any of the inputs This is useful for wrapping any tf.function that doesn't take a DTensor as input but would like to produce DTensor as result. The scope will also make sure all small constants are replicated as DTensors.

Args
`mesh`	A Mesh instance to extract a default mesh from.

Yields
A context in which all ops and tf.functions will run on the given mesh.

Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. For details, see the Google Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates. Some content is licensed under the numpy license.

Last updated 2024-04-26 UTC.

tf.experimental.dtensor.default_mesh

Args

Yields