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.
Yields | |
---|---|
A context in which all ops and tf.functions will run on the given mesh. |