Zachary Garrett
authored
Instead insert a `tf.function` decorator before (underneath) the `tff.tf_computation` decorator, and enter a `tf.init_scope` context during `tf.Variable` creation. The `tf.init_scope` hoists variable creation out of the `tf.function` to the nearest graph scope (provided by `tff.tf_comptuation`), which would otherwise result in an error from the the `tf.function` wrapper. More details on `tf.Variable` in `tf.function`: https://www.tensorflow.org/guide/function#variables Longer term, it may be nice to fold this automatic wrapping into the `tff.tf_computation` decorator, or re-introduce as `tff.tf2_computation`, however this requires all users to know they must use `tf.init_scope` during variable creation. Inside TFF packages (e.g. `tff.learning`) we can handle this for them, but its an extra point of concern for custom algorithm writers. PiperOrigin-RevId: 337935371
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
Name | Last commit | Last update |
---|---|---|
.. | ||
proto | ||
python | ||
tools | ||
BUILD | ||
__init__.py | ||
version.py |