Call auto_composite_tensor in AutoCompositeTensorBijector's metaclass, eliminating the need for the @auto_composite_tensor decorator on bijector subclasses.

PiperOrigin-RevId: 376966154

Author: emilyfertig

tensorflow_probability/python/bijectors/BUILD

@@ -124,7 +124,6 @@ multi_substrate_py_library(
     ],
     deps = [
         # numpy dep,
-        # six dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
         "//tensorflow_probability/python/internal:auto_composite_tensor",

tensorflow_probability/python/bijectors/bijector.py

@@ -24,7 +24,6 @@
 
 # Dependency imports
 import numpy as np
-import six
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python.internal import assert_util
@@ -91,8 +90,8 @@ def unflatten(info, xs):
     tree_util.register_pytree_node(cls, flatten, unflatten)
 
 
-@six.add_metaclass(_BijectorMeta)
-class Bijector(tf.Module):
+# TODO(emilyaf): Look at using `__init_subclass__` instead of a metaclass.
+class Bijector(tf.Module, metaclass=_BijectorMeta):
   r"""Interface for transformations of a `Distribution` sample.
 
   Bijectors can be used to represent any differentiable and injective
@@ -1779,20 +1778,36 @@ def _composite_tensor_shape_params(self):
       return ()
 
 
+auto_composite_tensor_bijector = functools.partial(
+    auto_composite_tensor.auto_composite_tensor,
+    omit_kwargs=('parameters',),
+    non_identifying_kwargs=('name',),
+    module_name='tfp.bijectors')
+
+
+class _AutoCompositeTensorBijectorMeta(_BijectorMeta):
+  """Metaclass for `AutoCompositeTensorBijector`."""
+
+  def __new__(mcs, classname, baseclasses, attrs):  # pylint: disable=bad-mcs-classmethod-argument
+    """Give subclasses their own type_spec, not an inherited one."""
+
+    cls = super(_AutoCompositeTensorBijectorMeta, mcs).__new__(  # pylint: disable=too-many-function-args
+        mcs, classname, baseclasses, attrs)
+    return auto_composite_tensor_bijector(cls)
+
+
 class AutoCompositeTensorBijector(
-    Bijector, auto_composite_tensor.AutoCompositeTensor):
+    Bijector, auto_composite_tensor.AutoCompositeTensor,
+    metaclass=_AutoCompositeTensorBijectorMeta):
   r"""Base for `CompositeTensor` bijectors with auto-generated `TypeSpec`s.
 
   `CompositeTensor` objects are able to pass in and out of `tf.function` and
   `tf.while_loop`, or serve as part of the signature of a TF saved model.
   `Bijector` subclasses that follow the contract of
   `tfp.experimental.auto_composite_tensor` may be defined as `CompositeTensor`s
-  by inheriting from `AutoCompositeTensorBijector` and applying a class
-  decorator as shown here:
+  by inheriting from `AutoCompositeTensorBijector`:
 
   ```python
-  @tfp.experimental.auto_composite_tensor(
-    omit_kwargs=('name',), module_name='my_module')
   class MyBijector(tfb.AutoCompositeTensorBijector):
 
     # The remainder of the subclass implementation is unchanged.
@@ -1801,13 +1816,6 @@ class MyBijector(tfb.AutoCompositeTensorBijector):
   pass
 
 
-auto_composite_tensor_bijector = functools.partial(
-    auto_composite_tensor.auto_composite_tensor,
-    omit_kwargs=('parameters',),
-    non_identifying_kwargs=('name',),
-    module_name='tfp.bijectors')
-
-
 def check_valid_ndims(ndims, validate=True):
   """Ensures that `ndims` is a non-negative integer.
 

tensorflow_probability/python/bijectors/bijector_test.py

@@ -27,7 +27,6 @@
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python import bijectors as tfb
-from tensorflow_probability.python.bijectors import bijector as bijector_lib
 from tensorflow_probability.python.internal import cache_util
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import prefer_static as ps
@@ -875,11 +874,10 @@ def testNestedCondition(self):
       mock_method.assert_called_once_with(mock.ANY, arg1=arg1, arg2=arg2)
 
 
-@bijector_lib.auto_composite_tensor_bijector
 class CompositeForwardBijector(tfb.AutoCompositeTensorBijector):
 
-  def __init__(self, scale=2., validate_args=False, name=None):
-    parameters = dict(locals())
+  def __init__(self, scale=2., validate_args=False, parameters=None, name=None):
+    parameters = dict(locals()) if parameters is None else parameters
     with tf.name_scope(name or 'forward_only') as name:
       self._scale = tensor_util.convert_nonref_to_tensor(
           scale,
@@ -897,6 +895,14 @@ def _forward_log_det_jacobian(self, _):
     return tf.math.log(self._scale)
 
 
+class CompositeForwardScaleThree(CompositeForwardBijector):
+
+  def __init__(self, name='scale_three'):
+    parameters = dict(locals())
+    super(CompositeForwardScaleThree, self).__init__(
+        scale=3., parameters=parameters, name=name)
+
+
 @test_util.test_all_tf_execution_regimes
 class AutoCompositeTensorBijectorTest(test_util.TestCase):
 
@@ -917,6 +923,15 @@ def test_disable_ct_bijector(self):
         non_ct_bijector.forward(x),
         tf.function(lambda b: b.forward(x))(unflat))
 
+  def test_composite_tensor_subclass(self):
+
+    bij = CompositeForwardScaleThree()
+    self.assertIs(bij._type_spec.value_type, type(bij))
+
+    flat = tf.nest.flatten(bij, expand_composites=True)
+    unflat = tf.nest.pack_sequence_as(bij, flat, expand_composites=True)
+    self.assertIsInstance(unflat, CompositeForwardScaleThree)
+
 
 if __name__ == '__main__':
   tf.test.main()

tensorflow_probability/python/internal/auto_composite_tensor.py

@@ -395,6 +395,32 @@ def _type_spec(self):
     pass
 
 
+def type_spec_register(name, allow_overwrite=True):
+  """Decorator used to register a unique name for a TypeSpec subclass.
+
+  Unlike TensorFlow's `type_spec.register`, this function allows a new
+  `TypeSpec` to be registered with a `name` that already appears in the
+  registry (overwriting the `TypeSpec` already registered with that name). This
+  allows for re-definition of `AutoCompositeTensor` subclasses in test
+  environments and iPython.
+
+  Args:
+    name: The name of the type spec. Must have the form
+    `"{project_name}.{type_name}"`.  E.g. `"my_project.MyTypeSpec"`.
+    allow_overwrite: `bool`, if `True` then the entry in the `TypeSpec` registry
+      keyed by `name` will be overwritten if it exists. If `False`, then
+      behavior is the same as `type_spec.register`.
+
+  Returns:
+    A class decorator that registers the decorated class with the given name.
+  """
+  # pylint: disable=protected-access
+  if allow_overwrite and name in type_spec._NAME_TO_TYPE_SPEC:
+    type_spec._TYPE_SPEC_TO_NAME.pop(
+        type_spec._NAME_TO_TYPE_SPEC.pop(name))
+  return type_spec.register(name)
+
+
 def auto_composite_tensor(
     cls=None, omit_kwargs=(), non_identifying_kwargs=(), module_name=None):
   """Automagically generate `CompositeTensor` behavior for `cls`.
@@ -540,19 +566,13 @@ def body(obj):
   type_spec_class_name = f'{cls.__name__}_ACTTypeSpec'
   type_spec_name = f'{module_name}.{type_spec_class_name}'
 
-  try:
-    ts = type_spec.lookup(type_spec_name)
-    return ts.value_type.fget(None)
-  except ValueError:
-    pass
-
   # If the declared class is already a CompositeTensor subclass, we can avoid
   # affecting the actual type of the returned class. Otherwise, we need to
   # explicitly mix in the CT type, and hence create and return a newly
   # synthesized type.
   if issubclass(cls, composite_tensor.CompositeTensor):
 
-    @type_spec.register(type_spec_name)
+    @type_spec_register(type_spec_name)
     class _AlreadyCTTypeSpec(_AutoCompositeTensorTypeSpec):
 
       @property
@@ -576,7 +596,6 @@ def _type_spec(obj):
   if clsid in _registry and issubclass(_registry[clsid], cls):
     return _registry[clsid]
 
-  @type_spec.register(type_spec_name)
   class _GeneratedCTTypeSpec(_AutoCompositeTensorTypeSpec):
 
     @property
@@ -595,4 +614,5 @@ def _type_spec(self):
 
   _AutoCompositeTensor.__name__ = cls.__name__
   _registry[clsid] = _AutoCompositeTensor
+  type_spec_register(type_spec_name)(_GeneratedCTTypeSpec)
   return _AutoCompositeTensor

tensorflow_probability/python/internal/backend/numpy/type_spec.py

@@ -14,13 +14,19 @@
 # ============================================================================
 """Numpy stub for `type_spec`."""
 
+import re
+
 __all__ = [
     'lookup',
     'register',
     'BatchableTypeSpec',
     'TypeSpec',
 ]
 
+_TYPE_SPEC_TO_NAME = {}
+_NAME_TO_TYPE_SPEC = {}
+_REGISTERED_NAME_RE = re.compile(r'^(\w+\.)+\w+$')
+
 
 def register(_):
   """No-op for registering a `tf.TypeSpec` for `saved_model`."""