Convert Chain, FillScaleTriL, and LambertWTransform bijectors to AutoCompositeTensor.

PiperOrigin-RevId: 378014572

Author: emilyfertig

tensorflow_probability/python/bijectors/bijector_properties_test.py

@@ -196,11 +196,6 @@
 })
 COMPOSITE_TENSOR_ATOL = collections.defaultdict(lambda: 1e-6)
 
-# TODO(b/182603117): Enable AutoCT for meta-bijectors.
-AUTO_COMPOSITE_TENSOR_IS_BROKEN = [
-    'FillScaleTriL',
-]
-
 
 def is_invert(bijector):
   return isinstance(bijector, (tfb.Invert, invert_lib._Invert))
@@ -915,24 +910,16 @@ def testCompositeTensor(self, bijector_name, data):
                         '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(
+      bijector = tf.nest.map_structure(
           lambda x: (tf.convert_to_tensor(x)  # pylint: disable=g-long-lambda
                      if isinstance(x, DeferredTensor) else x),
-          composite_bij,
+          bijector,
           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)
+    self.assertIsInstance(bijector, tf.__internal__.CompositeTensor)
+    flat = tf.nest.flatten(bijector, expand_composites=True)
+    unflat = tf.nest.pack_sequence_as(bijector, flat, expand_composites=True)
 
     # Compare forward maps before and after compositing.
     n = 3
@@ -950,15 +937,15 @@ def testCompositeTensor(self, bijector_name, data):
     # Input to tf.function
     self.assertAllClose(
         before_ys,
-        tf.function(lambda b: b.forward(xs))(composite_bij),
+        tf.function(lambda b: b.forward(xs))(bijector),
         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
+    wrt_vars = [xs] + [v for v in bijector.trainable_variables
                        if v.dtype.is_floating]
     with tf.GradientTape() as tape:
       tape.watch(wrt_vars)

tensorflow_probability/python/bijectors/chain.py

@@ -20,6 +20,7 @@
 from __future__ import print_function
 
 import tensorflow.compat.v2 as tf
+from tensorflow_probability.python.bijectors import bijector as bijector_lib
 from tensorflow_probability.python.bijectors import composition
 from tensorflow_probability.python.bijectors import ldj_ratio
 from tensorflow_probability.python.internal import parameter_properties
@@ -31,7 +32,7 @@
 ]
 
 
-class Chain(composition.Composition):
+class _Chain(composition.Composition):
   """Bijector which applies a sequence of bijectors.
 
   Example Use:
@@ -124,7 +125,7 @@ def __init__(self,
       inverse_min_event_ndims = None  # Inferred by base class.
 
     with tf.name_scope(name) as name:
-      super(Chain, self).__init__(
+      super(_Chain, self).__init__(
           bijectors=bijectors or (),
           validate_args=validate_args,
           validate_event_size=validate_event_size,
@@ -160,7 +161,32 @@ def _walk_inverse(self, step_fn, y, **kwargs):
     return y  # Now `x`
 
 
-@ldj_ratio.RegisterFLDJRatio(Chain)
+class Chain(_Chain, bijector_lib.AutoCompositeTensorBijector):
+
+  def __new__(cls, *args, **kwargs):
+    """Returns a `_Chain` if any of `bijectors` is not a `CompositeTensor."""
+    if cls is Chain:
+      if args:
+        bijectors = args[0]
+      else:
+        bijectors = kwargs.get('bijectors')
+
+      if bijectors is not None:
+        if not all(isinstance(b, tf.__internal__.CompositeTensor)
+                   for b in bijectors):
+          return _Chain(*args, **kwargs)
+    return super(Chain, cls).__new__(cls)
+
+
+Chain.__doc__ = _Chain.__doc__ + '\n' + (
+    'If every element of the `bijectors` list is a `CompositeTensor`, the '
+    'resulting `Chain` bijector is a `CompositeTensor` as well. If any element '
+    'of `bijectors` is not a `CompositeTensor`, then a non-`CompositeTensor` '
+    '`_Chain` instance is created instead. Bijector subclasses that inherit '
+    'from `Chain` will also inherit from `CompositeTensor`.')
+
+
+@ldj_ratio.RegisterFLDJRatio(_Chain)
 def _fldj_ratio_chain(p, x, q, y):
   """Sum-of-diffs FLDJRatio for Chains."""
   if len(p.bijectors) != len(q.bijectors):
@@ -177,7 +203,7 @@ def _fldj_ratio_chain(p, x, q, y):
   return tf.add_n(ratios)
 
 
-@ldj_ratio.RegisterILDJRatio(Chain)
+@ldj_ratio.RegisterILDJRatio(_Chain)
 def _ildj_ratio_chain(p, x, q, y):
   """Sum-of-diffs ILDJRatio for Chains."""
   if len(p.bijectors) != len(q.bijectors):

tensorflow_probability/python/bijectors/chain_test.py

@@ -404,6 +404,60 @@ def testDofChangeError(self):
     ignore_bij = tfb.Chain([exp, smc], validate_event_size=False)
     self.evaluate(ignore_bij.forward_log_det_jacobian([1., 2., 3.], 1))
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, disable_jax=True,
+      reason="Numpy and JAX have no notion of CompositeTensor/saved_model.")
+  def testCompositeTensor(self):
+    exp = tfb.Exp()
+    sp = tfb.Softplus()
+    aff = tfb.Scale(scale=2.)
+    chain = tfb.Chain(bijectors=[exp, sp, aff])
+    self.assertIsInstance(chain, tf.__internal__.CompositeTensor)
+
+    # Bijector may be flattened into `Tensor` components and rebuilt.
+    flat = tf.nest.flatten(chain, expand_composites=True)
+    unflat = tf.nest.pack_sequence_as(chain, flat, expand_composites=True)
+    self.assertIsInstance(unflat, tfb.Chain)
+
+    # Bijector may be input to a `tf.function`-decorated callable.
+    @tf.function
+    def call_forward(bij, x):
+      return bij.forward(x)
+
+    x = tf.ones([2, 3], dtype=tf.float32)
+    self.assertAllClose(call_forward(unflat, x), chain.forward(x))
+
+    # TypeSpec can be encoded/decoded.
+    struct_coder = tf.__internal__.saved_model.StructureCoder()
+    enc = struct_coder.encode_structure(chain._type_spec)
+    dec = struct_coder.decode_proto(enc)
+    self.assertEqual(chain._type_spec, dec)
+
+  def testNonCompositeTensor(self):
+
+    class NonCompositeScale(tfb.Bijector):
+      """Bijector that is not a `CompositeTensor`."""
+
+      def __init__(self, scale):
+        parameters = dict(locals())
+        self.scale = scale
+        super(NonCompositeScale, self).__init__(
+            validate_args=True,
+            forward_min_event_ndims=0.,
+            parameters=parameters,
+            name="non_composite_scale")
+
+      def _forward(self, x):
+        return x * self.scale
+
+      def _inverse(self, y):
+        return y / self.scale
+
+    exp = tfb.Exp()
+    scale = NonCompositeScale(scale=tf.constant(3.))
+    chain = tfb.Chain(bijectors=[exp, scale])
+    self.assertNotIsInstance(chain, tf.__internal__.CompositeTensor)
+    self.assertAllClose(chain.forward([1.]), exp.forward(scale.forward([1.])))
 
 if __name__ == "__main__":
   tf.test.main()

tensorflow_probability/python/bijectors/fill_scale_tril.py

@@ -33,7 +33,6 @@
 ]
 
 
-# TODO(b/182603117): Enable AutoCompositeTensor once Chain subclasses it.
 class FillScaleTriL(chain.Chain):
   """Transforms unconstrained vectors to TriL matrices with positive diagonal.
 
@@ -91,7 +90,8 @@ def __init__(self,
 
     Args:
       diag_bijector: `Bijector` instance, used to transform the output diagonal
-        to be positive.
+        to be positive. Must be an instance of `tf.__internal__.CompositeTensor`
+        (including `tfb.AutoCompositeTensorBijector`).
         Default value: `None` (i.e., `tfb.Softplus()`).
       diag_shift: Float value broadcastable and added to all diagonal entries
         after applying the `diag_bijector`. Setting a positive
@@ -104,11 +104,18 @@ def __init__(self,
         Default value: `False` (i.e., arguments are not validated).
       name: Python `str` name given to ops managed by this object.
         Default value: `fill_scale_tril`.
+
+    Raises:
+      TypeError, if `diag_bijector` is not an instance of
+        `tf.__internal__.CompositeTensor`.
     """
     parameters = dict(locals())
     with tf.name_scope(name) as name:
       if diag_bijector is None:
         diag_bijector = softplus.Softplus(validate_args=validate_args)
+      if not isinstance(diag_bijector, tf.__internal__.CompositeTensor):
+        raise TypeError('`diag_bijector` must be an instance of '
+                        '`tf.__internal__.CompositeTensor`.')
 
       if diag_shift is not None:
         dtype = dtype_util.common_dtype([diag_bijector, diag_shift], tf.float32)

tensorflow_probability/python/bijectors/lambertw_transform.py

@@ -106,7 +106,7 @@ def _w_delta_squared(z, delta):
 
 
 # Private class that implements the heavy tail transformation.
-class _HeavyTailOnly(bijector.Bijector):
+class _HeavyTailOnly(bijector.AutoCompositeTensorBijector):
   """Heavy tail transformation for Lambert W x F distributions.
 
   This bijector defines the transformation z = u * exp(0.5 * delta * u**2)
@@ -178,7 +178,6 @@ def _inverse_log_det_jacobian(self, y):
 # fix batch_shape inconsistencies when running distribution_properties_test.
 
 
-#  TODO(b/182603117): Enable AutoCompositeTensor when Chain is enabled.
 class LambertWTail(chain.Chain):
   """LambertWTail transformation for heavy-tail Lambert W x F random variables.