Skip to content
Snippets Groups Projects
user avatar
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
0e7c33dc
History
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