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.