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.
[null,null,["Last updated 2024-04-26 UTC."],[],[],null,["# tf.experimental.dtensor.run_on\n\n\u003cbr /\u003e\n\n|-------------------------------------------------------------------------------------------------------------------------|\n| [View source on GitHub](https://github.com/tensorflow/tensorflow/blob/v2.16.1/tensorflow/dtensor/python/api.py#L67-L86) |\n\nRuns enclosed functions in the DTensor device scope. (deprecated) \n\n @contextlib.contextmanager\n tf.experimental.dtensor.run_on(\n mesh: ../../../tf/experimental/dtensor/Mesh\n )\n\n| **Deprecated:** THIS FUNCTION IS DEPRECATED. It will be removed in a future version. Instructions for updating: Use [`dtensor.default_mesh`](../../../tf/experimental/dtensor/default_mesh) scope instead.\n\nThis function returns a scope. All the ops and tf.functions in this scope will\nrun on the DTensor device using the mesh provided.\nThis is useful for wrapping any tf.function that doesn't take a DTensor as\ninput but would like to produce DTensor as result. The scope will also make\nsure all small constants be replicated as DTensor.\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Args ---- ||\n|--------|-------------------------------------------------|\n| `mesh` | A Mesh instance to extract a default mesh from. |\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n\u003cbr /\u003e\n\n| Yields ------ ||\n|---|---|\n| A context in which all ops and tf.functions will run on the DTensor device. ||\n\n\u003cbr /\u003e"]]
View source on GitHub