Convert JointMap bijector to AutoCompositeTensor.

PiperOrigin-RevId: 377622021

Author: emilyfertig

tensorflow_probability/python/bijectors/joint_map.py

@@ -19,6 +19,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.python.util import nest  # pylint: disable=g-direct-tensorflow-import
 
@@ -28,7 +29,7 @@
 ]
 
 
-class JointMap(composition.Composition):
+class _JointMap(composition.Composition):
   """Bijector which applies a structure of bijectors in parallel.
 
   This is the "structured" counterpart to `Chain`. Whereas `Chain` applies an
@@ -92,7 +93,7 @@ def __init__(self,
       self._nested_structure = self._no_dependency(
           nest.map_structure(lambda b: None, bijectors))
 
-      super(JointMap, self).__init__(
+      super(_JointMap, self).__init__(
           bijectors=bijectors,
           validate_args=validate_args,
           forward_min_event_ndims=self._nested_structure,
@@ -115,3 +116,27 @@ def _walk_inverse(self, step_fn, ys, **kwargs):
         self._nested_structure,
         lambda bij, y: step_fn(bij, y, **kwargs.get(bij.name, {})),  # pylint: disable=unnecessary-lambda
         self._bijectors, ys, check_types=False)
+
+
+class JointMap(_JointMap, bijector_lib.AutoCompositeTensorBijector):
+
+  def __new__(cls, *args, **kwargs):
+    """Returns a `_JointMap` any of `bijectors` is not a `CompositeTensor."""
+    if cls is JointMap:
+      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 tf.nest.flatten(bijectors)):
+          return _JointMap(*args, **kwargs)
+    return super(JointMap, cls).__new__(cls)
+
+
+JointMap.__doc__ = _JointMap.__doc__ + '\n' + (
+    'If every element of `bijectors` is a `CompositeTensor`, the resulting '
+    '`JointMap` bijector is a `CompositeTensor` as well. If any element of '
+    '`bijectors` is not a `CompositeTensor`, then a non-`CompositeTensor` '
+    '`_JointMap` instance is created instead. Bijector subclasses that inherit '
+    'from `JointMap` will also inherit from `CompositeTensor`.')

tensorflow_probability/python/bijectors/joint_map_test.py

@@ -129,6 +129,61 @@ def test_inverse_has_event_ndims(self):
     bij_reshape.inverse_event_ndims([10])  # expect [9]
     self.assertEqual(bij_reshape.inverse_event_ndims([10]), [9])
 
+  @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.)
+    bij = tfb.JointMap(bijectors=[exp, sp, aff])
+    self.assertIsInstance(bij, tf.__internal__.CompositeTensor)
+
+    # Bijector may be flattened into `Tensor` components and rebuilt.
+    flat = tf.nest.flatten(bij, expand_composites=True)
+    unflat = tf.nest.pack_sequence_as(bij, flat, expand_composites=True)
+    self.assertIsInstance(unflat, tfb.JointMap)
+
+    # Bijector may be input to a `tf.function`-decorated callable.
+    @tf.function
+    def call_forward(bij, x):
+      return bij.forward(x)
+
+    x = [1., 2., 3.]
+    self.assertAllClose(call_forward(unflat, x), bij.forward(x))
+
+    # Type spec can be encoded/decoded.
+    struct_coder = tf.__internal__.saved_model.StructureCoder()
+    enc = struct_coder.encode_structure(bij._type_spec)
+    dec = struct_coder.decode_proto(enc)
+    self.assertEqual(bij._type_spec, dec)
+
+  def testNonCompositeTensor(self):
+
+    # TODO(b/182603117): Move NonComposite* into test_util.
+    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.))
+    bij = tfb.JointMap(bijectors=[exp, scale])
+    self.assertNotIsInstance(bij, tf.__internal__.CompositeTensor)
+    self.assertAllClose(
+        bij.forward([1., 1.]),
+        [exp.forward(1.), scale.forward(1.)])
+
 
 if __name__ == '__main__':
   tf.test.main()