tf.experimental.dtensor.run_on

Runs enclosed functions in the DTensor device scope.

@contextlib.contextmanager
tf.experimental.dtensor.run_on(
    layout_or_mesh: Union[tf.experimental.dtensor.Layout, tf.experimental.dtensor.Mesh]
)

This function returns a scope. All the ops and tf.functions in this scope will run on the DTensor device using the mesh provided or attached to the layout. 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 be replicated as DTensor.

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

Yields
A context in which all ops and tf.functions will run on the DTensor device.

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 2023-03-17 UTC.

tf.experimental.dtensor.run_on

Args

Yields