tf.experimental.dtensor.run_on
Stay organized with collections
Save and categorize content based on your preferences.
Runs enclosed functions in the DTensor device scope.
@contextlib.contextmanager
tf.experimental.dtensor.run_on(
mesh: 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.
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.
|
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-10-06 UTC.
[[["Easy to understand","easyToUnderstand","thumb-up"],["Solved my problem","solvedMyProblem","thumb-up"],["Other","otherUp","thumb-up"]],[["Missing the information I need","missingTheInformationINeed","thumb-down"],["Too complicated / too many steps","tooComplicatedTooManySteps","thumb-down"],["Out of date","outOfDate","thumb-down"],["Samples / code issue","samplesCodeIssue","thumb-down"],["Other","otherDown","thumb-down"]],["Last updated 2023-10-06 UTC."],[],[]]
View source on GitHub