Add non_identifying_kwargs to AutoCompositeTensor. These kwargs are preserved through flatten/unflatten and appear in serializations, but are omitted from comparison/equality checks and hashes.

emilyfertig · brianwa84 · commit 35932a7f9012 · 2021-05-20T13:45:40.000-04:00
PiperOrigin-RevId: 374718734
diff --git a/tensorflow_probability/python/internal/auto_composite_tensor.py b/tensorflow_probability/python/internal/auto_composite_tensor.py
@@ -60,7 +60,7 @@ def _deferred_assertion_context(is_deferred=True):
 
 _SENTINEL = object()
 
-_AUTO_COMPOSITE_TENSOR_VERSION = 2
+_AUTO_COMPOSITE_TENSOR_VERSION = 3
 
 # Cache maps __init__ method to signature
 _sig_cache = {}
@@ -171,7 +171,8 @@ class _AutoCompositeTensorTypeSpec(tf.TypeSpec):
                '_comparable')
 
   def __init__(self, param_specs, non_tensor_params, omit_kwargs,
-               prefer_static_value, callable_params=None):
+               prefer_static_value, non_identifying_kwargs,
+               callable_params=None):
     """Initializes a new `_AutoCompositeTensorTypeSpec`.
 
     Args:
@@ -189,6 +190,11 @@ def __init__(self, param_specs, non_tensor_params, omit_kwargs,
         of `Tensor`-like kwargs to the `AutoCompositeTensor`s constructor that
         may be stored as static values, if known. These are typically shapes or
         axis values.
+      non_identifying_kwargs: Python `tuple` of strings corresponding to the
+        names of kwargs to the `AutoCompositeTensor`s constructor whose values
+        are not relevant to the unique identification of the
+        `_AutoCompositeTensorTypeSpec` instance. Equality/comparison checks and
+        `__hash__` do not depend on these kwargs.
       callable_params: Python `dict` of callable kwargs to the
         `AutoCompositeTensor`'s constructor that do not subclass
         `CompositeTensor`, or `None`. If `callable_params` is a non-empty
@@ -199,22 +205,32 @@ def __init__(self, param_specs, non_tensor_params, omit_kwargs,
     self._non_tensor_params = non_tensor_params
     self._omit_kwargs = omit_kwargs
     self._prefer_static_value = prefer_static_value
+    self._non_identifying_kwargs = non_identifying_kwargs
     self._callable_params = {} if callable_params is None else callable_params
 
     self._serializable = (
         _AUTO_COMPOSITE_TENSOR_VERSION,
         self._param_specs,
         self._non_tensor_params,
         self._omit_kwargs,
-        self._prefer_static_value)
+        self._prefer_static_value,
+        self._non_identifying_kwargs)
 
-    # TODO(b/182603117): Distinguish between `omit_kwargs_from_constructor`
-    # and `omit_kwargs_for_comparison`.
-    self._comparable = self._serializable + (
-        tf.nest.map_structure(id, self._callable_params),)
+    def remove_kwargs(d):
+      return {k: v for k, v in d.items()
+              if k not in self._non_identifying_kwargs}
+
+    self._comparable = (
+        _AUTO_COMPOSITE_TENSOR_VERSION,
+        remove_kwargs(self._param_specs),
+        remove_kwargs(self._non_tensor_params),
+        self._omit_kwargs,
+        self._prefer_static_value,
+        self._non_identifying_kwargs,
+        tf.nest.map_structure(id, remove_kwargs(self._callable_params)))
 
   @classmethod
-  def from_instance(cls, instance, omit_kwargs=()):
+  def from_instance(cls, instance, omit_kwargs=(), non_identifying_kwargs=()):
     cls_value_type = cls.value_type.fget(None)
     if type(instance) is not cls_value_type:  # pylint: disable=unidiomatic-typecheck
       raise ValueError(f'`{type(instance).__name__}` has inherited the '
@@ -245,6 +261,7 @@ def from_instance(cls, instance, omit_kwargs=()):
                non_tensor_params=non_tensor_params,
                omit_kwargs=omit_kwargs,
                prefer_static_value=prefer_static_value,
+               non_identifying_kwargs=non_identifying_kwargs,
                callable_params=callable_params)
 
   def _to_components(self, obj):
@@ -273,6 +290,9 @@ def _deserialize(cls, encoded):
     if version == 1:
       encoded = encoded + ((),)
       version = 2
+    if version == 2:
+      encoded = encoded + ((),)
+      version = 3
     if version != _AUTO_COMPOSITE_TENSOR_VERSION:
       raise ValueError(f'Expected version {_AUTO_COMPOSITE_TENSOR_VERSION},'
                        f' but got {version}.')
@@ -375,7 +395,8 @@ def _type_spec(self):
     pass
 
 
-def auto_composite_tensor(cls=None, omit_kwargs=(), module_name=None):
+def auto_composite_tensor(
+    cls=None, omit_kwargs=(), non_identifying_kwargs=(), module_name=None):
   """Automagically generate `CompositeTensor` behavior for `cls`.
 
   `CompositeTensor` objects are able to pass in and out of `tf.function` and
@@ -499,6 +520,8 @@ def body(obj):
   Args:
     cls: The class for which to create a CompositeTensor subclass.
     omit_kwargs: Optional sequence of kwarg names to be omitted from the spec.
+    non_identifying_kwargs: Optional sequence of kwarg names to be omitted from
+      equality/comparison checks and the `__hash__` method of the spec.
     module_name: The module name with which to register the `TypeSpec`. If
       `None`, defaults to `cls.__module__`.
 
@@ -508,6 +531,7 @@ def body(obj):
   if cls is None:
     return functools.partial(auto_composite_tensor,
                              omit_kwargs=omit_kwargs,
+                             non_identifying_kwargs=non_identifying_kwargs,
                              module_name=module_name)
 
   if module_name is None:
@@ -537,11 +561,15 @@ def value_type(self):
 
     _AlreadyCTTypeSpec.__name__ = type_spec_class_name
 
-    cls._type_spec = property(  # pylint: disable=protected-access
-        lambda self: _AlreadyCTTypeSpec.from_instance(self, omit_kwargs))
+    def _type_spec(obj):
+      return _AlreadyCTTypeSpec.from_instance(
+          obj, omit_kwargs, non_identifying_kwargs)
+
+    cls._type_spec = property(_type_spec)  # pylint: disable=protected-access
     return cls
 
-  clsid = (cls.__module__, cls.__name__, omit_kwargs)
+  clsid = (cls.__module__, cls.__name__, omit_kwargs,
+           non_identifying_kwargs)
 
   # Check for subclass if retrieving from the _registry, in case the user
   # has redefined the class (e.g. in a REPL/notebook).
@@ -562,7 +590,8 @@ class _AutoCompositeTensor(cls, composite_tensor.CompositeTensor):
 
     @property
     def _type_spec(self):
-      return _GeneratedCTTypeSpec.from_instance(self, omit_kwargs)
+      return _GeneratedCTTypeSpec.from_instance(
+          self, omit_kwargs, non_identifying_kwargs)
 
   _AutoCompositeTensor.__name__ = cls.__name__
   _registry[clsid] = _AutoCompositeTensor
diff --git a/tensorflow_probability/python/internal/auto_composite_tensor_test.py b/tensorflow_probability/python/internal/auto_composite_tensor_test.py
@@ -49,20 +49,20 @@
 
 
 AutoIdentity = tfp.experimental.auto_composite_tensor(
-    tf.linalg.LinearOperatorIdentity, omit_kwargs=('name',))
+    tf.linalg.LinearOperatorIdentity, non_identifying_kwargs=('name',))
 AutoDiag = tfp.experimental.auto_composite_tensor(
-    tf.linalg.LinearOperatorDiag, omit_kwargs=('name',))
+    tf.linalg.LinearOperatorDiag, non_identifying_kwargs=('name',))
 AutoBlockDiag = tfp.experimental.auto_composite_tensor(
-    tf.linalg.LinearOperatorBlockDiag, omit_kwargs=('name',))
+    tf.linalg.LinearOperatorBlockDiag, non_identifying_kwargs=('name',))
 AutoTriL = tfp.experimental.auto_composite_tensor(
-    tf.linalg.LinearOperatorLowerTriangular, omit_kwargs=('name',))
+    tf.linalg.LinearOperatorLowerTriangular, non_identifying_kwargs=('name',))
 
 AutoNormal = tfp.experimental.auto_composite_tensor(
-    tfd.Normal, omit_kwargs=('name',))
+    tfd.Normal, non_identifying_kwargs=('name',))
 AutoIndependent = tfp.experimental.auto_composite_tensor(
-    tfd.Independent, omit_kwargs=('name',))
+    tfd.Independent, non_identifying_kwargs=('name',))
 AutoReshape = tfp.experimental.auto_composite_tensor(
-    tfb.Reshape, omit_kwargs=('name',))
+    tfb.Reshape, non_identifying_kwargs=('name',))
 
 
 class Model(tf.Module):
@@ -105,7 +105,7 @@ def tearDownModule():
 class AutoCompositeTensorTest(test_util.TestCase):
 
   def test_example(self):
-    @tfp.experimental.auto_composite_tensor(omit_kwargs=('name',))
+    @tfp.experimental.auto_composite_tensor(non_identifying_kwargs=('name',))
     class Adder(object):
 
       def __init__(self, x, y, name=None):
@@ -185,7 +185,7 @@ def test_preconditioner(self):
     tfed = tfp.experimental.distributions
     auto_ct_mvn_prec_linop = tfp.experimental.auto_composite_tensor(
         tfed.MultivariateNormalPrecisionFactorLinearOperator,
-        omit_kwargs=('name',))
+        non_identifying_kwargs=('name',))
     tril = AutoTriL(**cov_linop.cholesky().parameters)
     momentum_distribution = auto_ct_mvn_prec_linop(precision_factor=tril)
     def body(d):
@@ -408,15 +408,26 @@ def __init__(self):
     d_ct = AutoStandardNormal()
     self.assertLen(tf.nest.flatten(d_ct, expand_composites=True), 0)
 
+  def test_names_preserved_through_flatten(self):
+
+    dist = AutoNormal(0., scale=3., name='ScaleThreeNormal')
+    flat = tf.nest.flatten(dist, expand_composites=True)
+    unflat = tf.nest.pack_sequence_as(dist, flat, expand_composites=True)
+    unflat_name = ('ScaleThreeNormal' if tf.executing_eagerly()
+                   else 'ScaleThreeNormal_1')
+    self.assertEqual(unflat.name, unflat_name)
+
 
 class _TestTypeSpec(auto_composite_tensor._AutoCompositeTensorTypeSpec):
 
   def __init__(self, param_specs, non_tensor_params=None, omit_kwargs=(),
-               prefer_static_value=(), callable_params=None):
+               prefer_static_value=(), non_identifying_kwargs=(),
+               callable_params=None):
     non_tensor_params = {} if non_tensor_params is None else non_tensor_params
     super(_TestTypeSpec, self).__init__(
         param_specs, non_tensor_params=non_tensor_params,
         omit_kwargs=omit_kwargs, prefer_static_value=prefer_static_value,
+        non_identifying_kwargs=non_identifying_kwargs,
         callable_params=callable_params)
 
   @property
@@ -452,7 +463,16 @@ class AutoCompositeTensorTypeSpecTest(test_util.TestCase):
                         'b': tfb.Scale(3.)._type_spec},
            omit_kwargs=('name', 'foo'),
            prefer_static_value=('a',),
-           callable_params={'f': tf.math.exp}))
+           callable_params={'f': tf.math.exp})),
+      ('DifferentNonIdentifyingKwargsValues',
+       _TestTypeSpec(
+           param_specs={'x': tf.TensorSpec([], tf.float64)},
+           non_tensor_params={'name': 'MyAutoCT'},
+           non_identifying_kwargs=('name')),
+       _TestTypeSpec(
+           param_specs={'x': tf.TensorSpec([], tf.float64)},
+           non_tensor_params={'name': 'OtherAutoCT'},
+           non_identifying_kwargs=('name'))),
       )
   def testEquality(self, v1, v2):
     # pylint: disable=g-generic-assert
@@ -480,7 +500,15 @@ def testEquality(self, v1, v2):
        _TestTypeSpec(
            param_specs={'a': tf.TensorSpec([3, None], tf.float32)},
            omit_kwargs=('name', 'foo'),
-           callable_params={'f': tf.math.sigmoid}))
+           callable_params={'f': tf.math.sigmoid})),
+      ('DifferentMetadata',
+       _TestTypeSpec(
+           param_specs={'a': tf.TensorSpec([3, 2], tf.float32)},
+           non_tensor_params={'validate_args': True},
+           non_identifying_kwargs=('name',)),
+       _TestTypeSpec(
+           param_specs={'a': tf.TensorSpec([3, None], tf.float32)},
+           non_tensor_params={'validate_args': True})),
       )
   def testInequality(self, v1, v2):
     # pylint: disable=g-generic-assert
@@ -512,7 +540,16 @@ def testInequality(self, v1, v2):
            param_specs={'a': tf.TensorSpec([3, None], tf.float32),
                         'b': tfb.Scale(3.)._type_spec},
            omit_kwargs=('name', 'foo'),
-           callable_params={'f': tf.math.exp}))
+           callable_params={'f': tf.math.exp})),
+      ('DifferentNonIdentifyingKwargsValues',
+       _TestTypeSpec(
+           param_specs={'x': tf.TensorSpec(None, tf.float64)},
+           non_tensor_params={'name': 'MyAutoCT'},
+           non_identifying_kwargs=('name')),
+       _TestTypeSpec(
+           param_specs={'x': tf.TensorSpec([], tf.float64)},
+           non_tensor_params={'name': 'OtherAutoCT'},
+           non_identifying_kwargs=('name'))),
       )
   def testIsCompatibleWith(self, v1, v2):
     self.assertTrue(v1.is_compatible_with(v2))
@@ -625,7 +662,7 @@ def testMostSpecificCompatibleTypeException(self, v1, v2):
       ('WithoutCallable',
        _TestTypeSpec(
            param_specs={'a': tf.TensorSpec([4, 2], tf.float32)},
-           omit_kwargs=('name',))),
+           omit_kwargs=('parameters',), non_identifying_kwargs=('name',))),
       ('WithCallable',
        _TestTypeSpec(
            param_specs={'a': tf.TensorSpec(None, tf.float32),
@@ -636,7 +673,8 @@ def testMostSpecificCompatibleTypeException(self, v1, v2):
   def testRepr(self, spec):
     spec_data = (auto_composite_tensor._AUTO_COMPOSITE_TENSOR_VERSION,
                  spec._param_specs, spec._non_tensor_params, spec._omit_kwargs,
-                 spec._prefer_static_value, spec._callable_params)
+                 spec._prefer_static_value, spec._non_identifying_kwargs,
+                 spec._callable_params)
     self.assertEqual(repr(spec), f'_TestTypeSpec{spec_data}')
 
 if __name__ == '__main__':
diff --git a/tensorflow_probability/python/util/deferred_tensor_test.py b/tensorflow_probability/python/util/deferred_tensor_test.py
@@ -621,9 +621,9 @@ def testRepr(self):
     expected = (
         "_TransformedVariableSpec(input_spec=TensorSpec(shape=(4, 2), "
         "dtype=tf.float32, name=None), "
-        "transform_or_spec=Sigmoid_ACTTypeSpec(2, {}, "
-        "{'low': None, 'high': None, 'validate_args': True}, ('name',), (), "
-        "{}), dtype=<dtype: 'float64'>, name=None)")
+        "transform_or_spec=Sigmoid_ACTTypeSpec(3, {}, {'low': None, 'high': "
+        "None, 'validate_args': True}, ('name',), (), (), {}), dtype=<dtype: "
+        "'float64'>, name=None)")
     self.assertEqual(repr(spec), expected)
 
 
