Implements input_tensorspec for SimpleEncoder.

This also changes detail in the use of tf.function, as retracing can be triggered multiple times in the context of different graphs, breaking the class in certain scenarios.

PiperOrigin-RevId: 255297843

Author: PraChetit

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

@@ -77,17 +77,15 @@ def __init__(self, encoder, tensorspec):
     # These dictionaries are filled inside of the initial_state_fn and encode_fn
     # methods, to be used in encode_fn and decode_fn methods, respectively.
     # Decorated by tf.function, their necessary side effects are realized during
-    # call to get_concrete_function(). Because of fixed input_signatures, these
-    # are traced only once. See the tf.function tutorial for more details on
-    # the tracing semantics.
+    # call to get_concrete_function().
     state_py_structure = {}
     encoded_py_structure = {}
 
     @tf.function
     def initial_state_fn():
       state = encoder.initial_state()
-      assert not state_py_structure  # This should be traced only once.
-      state_py_structure['state'] = nest.map_structure(lambda _: None, state)
+      if not state_py_structure:
+        state_py_structure['state'] = nest.map_structure(lambda _: None, state)
       # Simplify the structure that needs to be manipulated by the user.
       return tuple(nest.flatten(state))
 
@@ -119,10 +117,10 @@ def encode_fn(x, flat_state):
       flat_encoded_py_structure, flat_encoded_tf_structure = (
           py_utils.split_dict_py_tf(flat_encoded_structure))
 
-      assert not encoded_py_structure  # This should be traced only once.
-      encoded_py_structure['full'] = nest.map_structure(lambda _: None,
-                                                        full_encoded_structure)
-      encoded_py_structure['flat_py'] = flat_encoded_py_structure
+      if not encoded_py_structure:
+        encoded_py_structure['full'] = nest.map_structure(
+            lambda _: None, full_encoded_structure)
+        encoded_py_structure['flat_py'] = flat_encoded_py_structure
       return flat_encoded_tf_structure, updated_flat_state
 
     @tf.function(input_signature=[
@@ -145,6 +143,12 @@ def decode_fn(encoded_structure):
     self._initial_state_fn = initial_state_fn
     self._encode_fn = encode_fn
     self._decode_fn = decode_fn
+    self._tensorspec = tensorspec
+
+  @property
+  def input_tensorspec(self):
+    """Returns `tf.TensorSpec` describing input expected by `SimpleEncoder`."""
+    return self._tensorspec
 
   def initial_state(self, name=None):
     """Returns the initial state.
@@ -182,8 +186,7 @@ def encode(self, x, state=None, name=None):
     """
     if state is None:
       state = self.initial_state()
-    with tf.name_scope(name, 'simple_encoder_encode',
-                       [x] + list(state)):
+    with tf.name_scope(name, 'simple_encoder_encode', [x] + list(state)):
       return self._encode_fn(x, state)
 
   def decode(self, encoded_x, name=None):
@@ -205,4 +208,3 @@ def decode(self, encoded_x, name=None):
     """
     with tf.name_scope(name, 'simple_encoder_decode', encoded_x.values()):
       return self._decode_fn(encoded_x)
-

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

@@ -202,6 +202,14 @@ def test_input_signature_enforced(self):
       bad_encoded_x.update({'x': x})
       encoder.decode(bad_encoded_x)
 
+  def test_input_tensorspec(self):
+    x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
+    encoder = simple_encoder.SimpleEncoder(
+        core_encoder.EncoderComposer(
+            test_utils.PlusOneOverNEncodingStage()).make(),
+        tf.TensorSpec.from_tensor(x))
+    self.assertTrue(encoder.input_tensorspec.is_compatible_with(x))
+
   @parameterized.parameters([1.0, 'str', object])
   def test_not_an_encoder_raises(self, not_an_encoder):
     """Tests invalid encoder argument."""