Convert Blockwise bijector to AutoCompositeTensor.

PiperOrigin-RevId: 375833034

Author: emilyfertig

tensorflow_probability/python/bijectors/BUILD

@@ -306,6 +306,7 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/internal:assert_util",
+        "//tensorflow_probability/python/internal:auto_composite_tensor",
         "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:tensorshape_util",
     ],

tensorflow_probability/python/bijectors/blockwise.py

@@ -21,11 +21,14 @@
 import numpy as np
 import tensorflow.compat.v2 as tf
 
+from tensorflow_probability.python.bijectors import bijector as bijector_lib
 from tensorflow_probability.python.bijectors import chain
+from tensorflow_probability.python.bijectors import composition
 from tensorflow_probability.python.bijectors import invert
 from tensorflow_probability.python.bijectors import joint_map
 from tensorflow_probability.python.bijectors import split
 from tensorflow_probability.python.internal import assert_util
+from tensorflow_probability.python.internal import auto_composite_tensor
 from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.internal import tensorshape_util
 
@@ -42,8 +45,7 @@ def _get_static_splits(splits):
   return splits if static_splits is None else static_splits
 
 
-# TODO(b/182603117): Enable AutoCompositeTensor once Chain subclasses it.
-class Blockwise(chain.Chain):
+class _Blockwise(composition.Composition):
   """Bijector which applies a list of bijectors to blocks of a `Tensor`.
 
   More specifically, given [F_0, F_1, ... F_n] which are scalar or vector
@@ -151,9 +153,12 @@ def __init__(self,
           name='concat')
 
       self._maybe_changes_size = maybe_changes_size
-      super(Blockwise, self).__init__(
-          bijectors=[b_concat, b_joint, b_split],
+      self._chain = chain.Chain(
+          [b_concat, b_joint, b_split], validate_args=validate_args)
+      super(_Blockwise, self).__init__(
+          bijectors=self._chain.bijectors,
           validate_args=validate_args,
+          validate_event_size=True,
           parameters=parameters,
           name=name)
 
@@ -186,13 +191,13 @@ def inverse_block_sizes(self):
     return self._b_concat.split_sizes
 
   def _forward(self, x, **kwargs):
-    y = super(Blockwise, self)._forward(x, **kwargs)
+    y = super(_Blockwise, self)._forward(x, **kwargs)
     if not self._maybe_changes_size:
       tensorshape_util.set_shape(y, x.shape)
     return y
 
   def _inverse(self, y, **kwargs):
-    x = super(Blockwise, self)._inverse(y, **kwargs)
+    x = super(_Blockwise, self)._inverse(y, **kwargs)
     if not self._maybe_changes_size:
       tensorshape_util.set_shape(x, y.shape)
     return x
@@ -220,19 +225,19 @@ def _inverse_event_shape(self, output_shape):
   def _forward_event_shape_tensor(self, x, **kwargs):
     if not self._maybe_changes_size:
       return x
-    return super(Blockwise, self)._forward_event_shape_tensor(x, **kwargs)
+    return super(_Blockwise, self)._forward_event_shape_tensor(x, **kwargs)
 
   def _inverse_event_shape_tensor(self, y, **kwargs):
     if not self._maybe_changes_size:
       return y
-    return super(Blockwise, self)._inverse_event_shape_tensor(y, **kwargs)
+    return super(_Blockwise, self)._inverse_event_shape_tensor(y, **kwargs)
 
   def _walk_forward(self, step_fn, x, **kwargs):
-    return super(Blockwise, self)._walk_forward(
+    return self._chain._walk_forward(  # pylint: disable=protected-access
         step_fn, x, **{self._b_joint.name: kwargs})
 
   def _walk_inverse(self, step_fn, x, **kwargs):
-    return super(Blockwise, self)._walk_inverse(
+    return self._chain._walk_inverse(  # pylint: disable=protected-access
         step_fn, x, **{self._b_joint.name: kwargs})
 
 
@@ -263,3 +268,32 @@ def _validate_block_sizes(block_sizes, bijectors, validate_args):
   # Set the shape if missing to pass statically known structure to split.
   tensorshape_util.set_shape(block_sizes, [len(bijectors)])
   return block_sizes
+
+
+@bijector_lib.auto_composite_tensor_bijector
+class Blockwise(_Blockwise, auto_composite_tensor.AutoCompositeTensor):
+
+  def __new__(cls, *args, **kwargs):
+    """Returns a `_Blockwise` if any of `bijectors` is not `CompositeTensor."""
+    if cls is Blockwise:
+      if args:
+        bijectors = args[0]
+      elif 'bijectors' in kwargs:
+        bijectors = kwargs['bijectors']
+      else:
+        raise TypeError(
+            '`Blockwise.__new__()` is missing argument `bijectors`.')
+
+      if not all(isinstance(b, tf.__internal__.CompositeTensor)
+                 for b in bijectors):
+        return _Blockwise(*args, **kwargs)
+    return super(Blockwise, cls).__new__(cls)
+
+
+Blockwise.__doc__ = _Blockwise.__doc__ + '\n' + (
+    'If every element of the `bijectors` list is a `CompositeTensor`, the '
+    'resulting `Blockwise` bijector is a `CompositeTensor` as well. If any '
+    'element of `bijectors` is not a `CompositeTensor`, then a '
+    'non-`CompositeTensor` `_Blockwise` instance is created instead. Bijector '
+    'subclasses that inherit from `Blockwise` will also inherit from '
+    '`CompositeTensor`.')

tensorflow_probability/python/bijectors/blockwise_test.py

@@ -294,6 +294,60 @@ def testKwargs(self):
     bijectors[0]._inverse_log_det_jacobian.assert_called_with(mock.ANY, arg=7)
     bijectors[1]._inverse_log_det_jacobian.assert_called_with(mock.ANY, arg=8)
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, disable_jax=True,
+      reason='Numpy and JAX have no notion of CompositeTensor.')
+  def testCompositeTensor(self):
+    exp = tfb.Exp()
+    sp = tfb.Softplus()
+    aff = tfb.Scale(scale=2.)
+    blockwise = tfb.Blockwise(bijectors=[exp, sp, aff])
+    self.assertIsInstance(blockwise, tf.__internal__.CompositeTensor)
+
+    # Bijector may be flattened into `Tensor` components and rebuilt.
+    flat = tf.nest.flatten(blockwise, expand_composites=True)
+    unflat = tf.nest.pack_sequence_as(blockwise, flat, expand_composites=True)
+    self.assertIsInstance(unflat, tfb.Blockwise)
+
+    # 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), blockwise.forward(x))
+
+    # Type spec can be encoded/decoded.
+    struct_coder = tf.__internal__.saved_model.StructureCoder()
+    enc = struct_coder.encode_structure(blockwise._type_spec)
+    dec = struct_coder.decode_proto(enc)
+    self.assertEqual(blockwise._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
+
+    exp = tfb.Exp()
+    scale = NonCompositeScale(scale=tf.constant(3.))
+    blockwise = tfb.Blockwise(bijectors=[exp, scale])
+    self.assertNotIsInstance(blockwise, tf.__internal__.CompositeTensor)
+    self.assertAllClose(
+        blockwise.forward([1., 1.]),
+        tf.convert_to_tensor([exp.forward(1.), scale.forward(1.)]))
+
 
 if __name__ == '__main__':
   tf.test.main()

tensorflow_probability/python/bijectors/glow_test.py

@@ -20,6 +20,7 @@
 import tensorflow.compat.v2 as tf
 import tensorflow_probability as tfp
 from tensorflow_probability.python.bijectors import bijector_test_util
+from tensorflow_probability.python.bijectors import blockwise
 from tensorflow_probability.python.bijectors import composition
 from tensorflow_probability.python.internal import test_util
 from tensorflow_probability.python.math.gradient import batch_jacobian
@@ -116,7 +117,7 @@ def testDataInit_inverse(self):
 
     x = bijection.bijectors[0].inverse(x)
     for b in bijection.bijectors[1].bijectors:
-      if isinstance(b, tfb.blockwise.Blockwise):
+      if isinstance(b, blockwise._Blockwise):
         x1, x2 = tf.split(x, splits[-2-nblocks], axis=-1)
 
         for bb in b.bijectors[0].bijectors:
@@ -157,7 +158,7 @@ def testDataInit_forward(self):
     splits = [[bs[0]+bs[1], bs[2]] for bs in splits]
 
     for b in reversed(bijection.bijectors[1].bijectors):
-      if isinstance(b, tfb.blockwise.Blockwise):
+      if isinstance(b, blockwise._Blockwise):
         y1, y2 = tf.split(y, splits[nblocks], axis=-1)
 
         for bb in reversed(b.bijectors[0].bijectors):