Remove tf.function to prevent frequent amounts of retracing.

Outer level scopes should apply `tf.function` if desired.

PiperOrigin-RevId: 447591780

Author: ZacharyGarrett

tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder.py

@@ -13,10 +13,6 @@
 # limitations under the License.
 """Base Encoder class for encoding in the "many-to-one" case."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import collections
 import tensorflow as tf
 
@@ -28,7 +24,7 @@
 _TENSORS = 'tensors'
 
 
-class GatherEncoder(object):
+class GatherEncoder:
   """A class for a gather-like operations with encoding.
 
   This class provides functionality for encoding in the "many-to-one" case,
@@ -198,17 +194,19 @@ def _add_to_py_values(key, value):
       if key not in internal_py_values:
         internal_py_values[key] = value
 
-    @tf.function
     def initial_state_fn():
       """See the `initial_state` method of this class."""
-      state = encoder.initial_state()
+      # Convert to tensor values here. If the initial state is returning
+      # variables, this captures the value of the variable at the moment of
+      # this call.
+      state = tf.nest.map_structure(tf.convert_to_tensor,
+                                    encoder.initial_state())
       _add_to_structure('state', state)
       return tuple(tf.nest.flatten(state))
 
     state = initial_state_fn()
     flat_state_spec = tf.nest.map_structure(tf.TensorSpec.from_tensor, state)
 
-    @tf.function
     def get_params_fn(flat_state):
       """See the `get_params` method of this class."""
       py_utils.assert_compatible(flat_state_spec, flat_state)
@@ -249,16 +247,15 @@ def get_params_fn(flat_state):
               tuple(tf.nest.flatten(decode_before_sum_params_tf)),
               tuple(tf.nest.flatten(decode_after_sum_params_tf)))
 
-    encode_params, decode_before_sum_params, decode_after_sum_params = (
-        get_params_fn(state))
+    encode_params, decode_before_sum_params, decode_after_sum_params = tf.nest.map_structure(
+        tf.convert_to_tensor, (get_params_fn(state)))
     encode_params_spec = tf.nest.map_structure(tf.TensorSpec.from_tensor,
                                                encode_params)
     decode_before_sum_params_spec = tf.nest.map_structure(
         tf.TensorSpec.from_tensor, decode_before_sum_params)
     decode_after_sum_params_spec = tf.nest.map_structure(
         tf.TensorSpec.from_tensor, decode_after_sum_params)
 
-    @tf.function
     def encode_fn(x, params):
       """See the `encode` method of this class."""
       if not tensorspec.is_compatible_with(x):
@@ -294,7 +291,6 @@ def encode_fn(x, params):
     encoded_structure_spec = tf.nest.map_structure(tf.TensorSpec.from_tensor,
                                                    encoded_structure)
 
-    @tf.function
     def decode_before_sum_fn(encoded_structure, params):
       """See the `decode_before_sum` method of this class."""
       py_utils.assert_compatible(encoded_structure_spec, encoded_structure)
@@ -326,7 +322,6 @@ def decode_before_sum_fn(encoded_structure, params):
     part_decoded_structure_spec = tf.nest.map_structure(
         tf.TensorSpec.from_tensor, part_decoded_structure)
 
-    @tf.function
     def decode_after_sum_fn(part_decoded_structure, params, num_summands):
       """See the `decode_after_sum` method of this class."""
       py_utils.assert_compatible(part_decoded_structure_spec,
@@ -350,7 +345,6 @@ def decode_after_sum_fn(part_decoded_structure, params, num_summands):
                                     decode_after_sum_params, 1)
     assert tensorspec.is_compatible_with(decoded_x)
 
-    @tf.function
     def update_state_fn(flat_state, state_update_tensors):
       """See the `update_state` method of this class."""
       py_utils.assert_compatible(flat_state_spec, flat_state)

tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder_test.py

@@ -12,10 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import collections
 
 from absl.testing import parameterized
@@ -136,8 +132,8 @@ def test_python_constants_not_exposed(self):
                 test_utils.PlusOneEncodingStage(), P1_VALS).add_parent(
                     test_utils.SimpleLinearEncodingStage(2.0, 3.0),
                     SL_VALS).make(), tensorspec)
-    a_var = tf.compat.v1.get_variable('a_var', initializer=2.0)
-    b_var = tf.compat.v1.get_variable('b_var', initializer=3.0)
+    a_var = tf.Variable(2.0, name='a_var')
+    b_var = tf.Variable(3.0, name='b_var')
     encoder_tf = gather_encoder.GatherEncoder.from_encoder(
         core_encoder.EncoderComposer(
             test_utils.SimpleLinearEncodingStage(a_var, b_var)).add_parent(