Skip to content
  • Zachary Garrett's avatar
    Stop creating additional python functions wrapped in `tf.function` decorators. · 0e7c33dc
    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