tf.experimental.dtensor.run_on

Runs enclosed functions in the DTensor device scope. (deprecated)

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

Deprecated: THIS FUNCTION IS DEPRECATED. It will be removed in a future version. Instructions for updating: Use dtensor.default_mesh scope instead.

This function returns a scope. All the ops and tf.functions in this scope will run on the DTensor device using the mesh provided. 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
mesh	A Mesh instance to extract a default mesh from.

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