Temporarily remove AutoCompositeTensor from bijectors for TFP 0.13 release.

PiperOrigin-RevId: 374929567

Author: emilyfertig

tensorflow_probability/python/bijectors/bijector.py

@@ -20,7 +20,7 @@
 
 import abc
 import contextlib
-# import functools
+import functools
 
 # Dependency imports
 import numpy as np
@@ -1599,39 +1599,33 @@ def _composite_tensor_shape_params(self):
       return ()
 
 
-# Temporarily disable AutoCT for TFP 0.13 release
-# class AutoCompositeTensorBijector(
-#     Bijector, auto_composite_tensor.AutoCompositeTensor):
-#   r"""Base for `CompositeTensor` bijectors with auto-generated `TypeSpec`s.
+class AutoCompositeTensorBijector(
+    Bijector, auto_composite_tensor.AutoCompositeTensor):
+  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  # pylint: disable=line-too-long
-#   by inheriting from `AutoCompositeTensorBijector` and applying a class
-#   decorator as shown here:
+  `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:
 
-#   ```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.
-#   ```
-#   """
-#   pass
-
-
-# auto_composite_tensor_bijector = functools.partial(
-#     auto_composite_tensor.auto_composite_tensor,
-#     omit_kwargs=('parameters',),
-#     non_identifying_kwargs=('name',),
-#     module_name='tfp.bijectors')
+  ```python
+  @tfp.experimental.auto_composite_tensor(
+    omit_kwargs=('name',), module_name='my_module')
+  class MyBijector(tfb.AutoCompositeTensorBijector):
 
-AutoCompositeTensorBijector = Bijector
+    # The remainder of the subclass implementation is unchanged.
+  ```
+  """
+  pass
 
 
-auto_composite_tensor_bijector = lambda cls, **kwargs: cls
+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):

tensorflow_probability/python/bijectors/bijector_properties_test.py

@@ -28,13 +28,14 @@
 from tensorflow_probability.python import bijectors as tfb
 from tensorflow_probability.python import experimental
 from tensorflow_probability.python.bijectors import hypothesis_testlib as bijector_hps
+from tensorflow_probability.python.bijectors import invert as invert_lib
 from tensorflow_probability.python.internal import hypothesis_testlib as tfp_hps
 from tensorflow_probability.python.internal import prefer_static
 from tensorflow_probability.python.internal import samplers
 from tensorflow_probability.python.internal import tensor_util
 from tensorflow_probability.python.internal import tensorshape_util
 from tensorflow_probability.python.internal import test_util
-# from tensorflow_probability.python.util.deferred_tensor import DeferredTensor
+from tensorflow_probability.python.util.deferred_tensor import DeferredTensor
 
 
 TF2_FRIENDLY_BIJECTORS = (
@@ -185,7 +186,6 @@
 
 COMPOSITE_TENSOR_IS_BROKEN = [
     'BatchNormalization',  # tf.layers arg
-    'Inline',  # callable
     'RationalQuadraticSpline',  # TODO(b/185628453): Debug loss of static info.
 ]
 
@@ -204,7 +204,7 @@
 
 
 def is_invert(bijector):
-  return isinstance(bijector, tfb.Invert)
+  return isinstance(bijector, (tfb.Invert, invert_lib._Invert))
 
 
 def is_transform_diagonal(bijector):
@@ -244,8 +244,8 @@ def _constraint(param):
 
 
 # TODO(b/141098791): Eliminate this.
-# @experimental.auto_composite_tensor
-class CallableModule(tf.Module):  # , experimental.AutoCompositeTensor):
+@experimental.auto_composite_tensor
+class CallableModule(tf.Module, experimental.AutoCompositeTensor):
   """Convenience object for capturing variables closed over by Inline."""
 
   def __init__(self, fn, varobj):
@@ -887,16 +887,38 @@ def testEquality(self, bijector_name, data):
   @hp.given(hps.data())
   @tfp_hps.tfp_hp_settings()
   def testCompositeTensor(self, bijector_name, data):
-    # Test that making a composite tensor of this bijector doesn't throw any
-    # errors.
+
     bijector, event_dim = self._draw_bijector(
-        bijector_name, data, batch_shape=[],
+        bijector_name, data,
+        batch_shape=[],
+        validate_args=True,
         allowed_bijectors=(set(TF2_FRIENDLY_BIJECTORS) -
                            set(COMPOSITE_TENSOR_IS_BROKEN)))
-    composite_bij = experimental.as_composite(bijector)
+
+    if type(bijector) is invert_lib._Invert:  # pylint: disable=unidiomatic-typecheck
+      if isinstance(bijector.bijector, tf.__internal__.CompositeTensor):
+        raise TypeError('`_Invert` should wrap only non-`CompositeTensor` '
+                        'bijectors.')
+      self.skipTest('`_Invert` bijectors are not `CompositeTensor`s.')
+
+    # TODO(b/182603117): Remove "if" condition and s/composite_bij/bijector
+    # when AutoCT is enabled for meta-bijectors and LinearOperator.
+    if type(bijector).__name__ in AUTO_COMPOSITE_TENSOR_IS_BROKEN:
+      composite_bij = experimental.as_composite(bijector)
+    else:
+      composite_bij = bijector
+
+    if not tf.executing_eagerly():
+      composite_bij = tf.nest.map_structure(
+          lambda x: (tf.convert_to_tensor(x)  # pylint: disable=g-long-lambda
+                     if isinstance(x, DeferredTensor) else x),
+          composite_bij,
+          expand_composites=True)
+
+    self.assertIsInstance(composite_bij, tf.__internal__.CompositeTensor)
     flat = tf.nest.flatten(composite_bij, expand_composites=True)
-    unflat = tf.nest.pack_sequence_as(composite_bij, flat,
-                                      expand_composites=True)
+    unflat = tf.nest.pack_sequence_as(
+        composite_bij, flat, expand_composites=True)
 
     # Compare forward maps before and after compositing.
     n = 3
@@ -911,6 +933,26 @@ def testCompositeTensor(self, bijector_name, data):
     after_xs = unflat.inverse(ys)
     self.assertAllClose(*self.evaluate((before_xs, after_xs)))
 
+    # Input to tf.function
+    self.assertAllClose(
+        before_ys,
+        tf.function(lambda b: b.forward(xs))(composite_bij),
+        rtol=COMPOSITE_TENSOR_RTOL[bijector_name],
+        atol=COMPOSITE_TENSOR_ATOL[bijector_name])
+
+    # Forward mapping: Check differentiation through forward mapping with
+    # respect to the input and parameter variables.  Also check that any
+    # variables are not referenced overmuch.
+    xs = self._draw_domain_tensor(bijector, data, event_dim)
+    wrt_vars = [xs] + [v for v in composite_bij.trainable_variables
+                       if v.dtype.is_floating]
+    with tf.GradientTape() as tape:
+      tape.watch(wrt_vars)
+      # TODO(b/73073515): Fix graph mode gradients with bijector caching.
+      ys = bijector.forward(xs + 0)
+    grads = tape.gradient(ys, wrt_vars)
+    assert_no_none_grad(bijector, 'forward', wrt_vars, grads)
+
 
 def ensure_nonzero(x):
   return tf.where(x < 1e-6, tf.constant(1e-3, x.dtype), x)

tensorflow_probability/python/bijectors/bijector_test.py

@@ -790,26 +790,25 @@ def _forward_log_det_jacobian(self, _):
     return tf.math.log(self._scale)
 
 
-  # Test disabled temporarily for TFP 0.13 release.
-# @test_util.test_all_tf_execution_regimes
-# class AutoCompositeTensorBijectorTest(test_util.TestCase):
+@test_util.test_all_tf_execution_regimes
+class AutoCompositeTensorBijectorTest(test_util.TestCase):
 
-#   def test_disable_ct_bijector(self):
+  def test_disable_ct_bijector(self):
 
-#     ct_bijector = CompositeForwardBijector()
-#     self.assertIsInstance(ct_bijector, tf.__internal__.CompositeTensor)
+    ct_bijector = CompositeForwardBijector()
+    self.assertIsInstance(ct_bijector, tf.__internal__.CompositeTensor)
 
-#     non_ct_bijector = ForwardOnlyBijector()
-#     self.assertNotIsInstance(non_ct_bijector, tf.__internal__.CompositeTensor)
+    non_ct_bijector = ForwardOnlyBijector()
+    self.assertNotIsInstance(non_ct_bijector, tf.__internal__.CompositeTensor)
 
-#     flat = tf.nest.flatten(ct_bijector, expand_composites=True)
-#     unflat = tf.nest.pack_sequence_as(
-#         ct_bijector, flat, expand_composites=True)
+    flat = tf.nest.flatten(ct_bijector, expand_composites=True)
+    unflat = tf.nest.pack_sequence_as(
+        ct_bijector, flat, expand_composites=True)
 
-#     x = tf.constant([2., 3.])
-#     self.assertAllClose(
-#         non_ct_bijector.forward(x),
-#         tf.function(lambda b: b.forward(x))(unflat))
+    x = tf.constant([2., 3.])
+    self.assertAllClose(
+        non_ct_bijector.forward(x),
+        tf.function(lambda b: b.forward(x))(unflat))
 
 
 if __name__ == '__main__':

tensorflow_probability/python/bijectors/exp.py

@@ -23,6 +23,7 @@
 from tensorflow_probability.python.bijectors import bijector as bijector_lib
 from tensorflow_probability.python.bijectors import invert
 from tensorflow_probability.python.bijectors import power_transform
+from tensorflow_probability.python.internal import auto_composite_tensor
 
 
 __all__ = [
@@ -76,7 +77,8 @@ def __init__(self,
 # TODO(b/182603117): Remove `AutoCompositeTensor` when `Invert` subclasses
 # `AutoCompositeTensor` and ensure `tf.saved_model` still works.
 @bijector_lib.auto_composite_tensor_bijector
-class Log(invert.Invert):
+class Log(invert.Invert,
+          auto_composite_tensor.AutoCompositeTensor):
   """Compute `Y = log(X)`. This is `Invert(Exp())`."""
 
   def __init__(self, validate_args=False, name='log'):

tensorflow_probability/python/bijectors/expm1.py

@@ -21,6 +21,7 @@
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import invert
+from tensorflow_probability.python.internal import auto_composite_tensor
 
 
 __all__ = [
@@ -94,7 +95,7 @@ def _forward_log_det_jacobian(self, x):
 # TODO(b/182603117): Remove `AutoCompositeTensor` when `Invert` subclasses
 # `AutoCompositeTensor`.
 @bijector.auto_composite_tensor_bijector
-class Log1p(invert.Invert):
+class Log1p(invert.Invert, auto_composite_tensor.AutoCompositeTensor):
   """Compute `Y = log1p(X)`. This is `Invert(Expm1())`."""
 
   def __init__(self, validate_args=False, name='log1p'):

tensorflow_probability/python/bijectors/invert.py

@@ -21,14 +21,14 @@
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python.bijectors import bijector as bijector_lib
-# from tensorflow_probability.python.internal import auto_composite_tensor
+from tensorflow_probability.python.internal import auto_composite_tensor
 
 __all__ = [
     'Invert',
 ]
 
 
-class Invert(bijector_lib.Bijector):
+class _Invert(bijector_lib.Bijector):
   """Bijector which inverts another Bijector.
 
   Example Use: [ExpGammaDistribution (see Background & Context)](
@@ -73,7 +73,7 @@ def __init__(self, bijector, validate_args=False, parameters=None, name=None):
     name = name or '_'.join(['invert', bijector.name])
     with tf.name_scope(name) as name:
       self._bijector = bijector
-      super(Invert, self).__init__(
+      super(_Invert, self).__init__(
           forward_min_event_ndims=bijector.inverse_min_event_ndims,
           inverse_min_event_ndims=bijector.forward_min_event_ndims,
           dtype=bijector.dtype,
@@ -138,28 +138,26 @@ def forward_event_ndims(self, event_ndims, **kwargs):
     return self.bijector.inverse_event_ndims(event_ndims, **kwargs)
 
 
-# Temporarily removing AutoCompositeTensor for TFP 0.13 release.
-# pylint: disable=line-too-long
-# @bijector_lib.auto_composite_tensor_bijector
-# class Invert(_Invert, auto_composite_tensor.AutoCompositeTensor):
+@bijector_lib.auto_composite_tensor_bijector
+class Invert(_Invert, auto_composite_tensor.AutoCompositeTensor):
 
-#   def __new__(cls, *args, **kwargs):
-#     """Returns an `_Invert` instance if `bijector` is not a `CompositeTensor."""
-#     if cls is Invert:
-#       if args:
-#         bijector = args[0]
-#       elif 'bijector' in kwargs:
-#         bijector = kwargs['bijector']
-#       else:
-#         raise TypeError('`Invert.__new__()` is missing argument `bijector`.')
+  def __new__(cls, *args, **kwargs):
+    """Returns an `_Invert` instance if `bijector` is not a `CompositeTensor."""
+    if cls is Invert:
+      if args:
+        bijector = args[0]
+      elif 'bijector' in kwargs:
+        bijector = kwargs['bijector']
+      else:
+        raise TypeError('`Invert.__new__()` is missing argument `bijector`.')
 
-#       if not isinstance(bijector, tf.__internal__.CompositeTensor):
-#         return _Invert(*args, **kwargs)
-#     return super(Invert, cls).__new__(cls)
+      if not isinstance(bijector, tf.__internal__.CompositeTensor):
+        return _Invert(*args, **kwargs)
+    return super(Invert, cls).__new__(cls)
 
 
-# Invert.__doc__ = _Invert.__doc__ + '/n' + (
-#     'When an `Invert` bijector is constructed, if its `bijector` arg is not a '
-#     '`CompositeTensor` instance, an `_Invert` instance is returned instead. '
-#     'Bijectors subclasses that inherit from `Invert` will also inherit from '
-#     ' `CompositeTensor`.')
+Invert.__doc__ = _Invert.__doc__ + '/n' + (
+    'When an `Invert` bijector is constructed, if its `bijector` arg is not a '
+    '`CompositeTensor` instance, an `_Invert` instance is returned instead. '
+    'Bijectors subclasses that inherit from `Invert` will also inherit from '
+    ' `CompositeTensor`.')

tensorflow_probability/python/experimental/composite_tensor_test.py

@@ -454,11 +454,7 @@ def test_basics_mixture_same_family(self):
     self.evaluate(unflat.log_prob(.5))
 
   def test_already_composite_tensor(self):
-    AutoScale = tfp.experimental.auto_composite_tensor(  # pylint: disable=invalid-name
-        tfb.Scale, omit_kwargs=('parameters',),
-        non_identifying_kwargs=('name',),
-        module_name=('tfp.bijectors'))
-    b = AutoScale(2.)
+    b = tfb.Scale(2.)
     b2 = tfp.experimental.as_composite(b)
     self.assertIsInstance(b, tf.__internal__.CompositeTensor)
     self.assertIs(b, b2)