Subclass Composition instead of Chain and Blockwise in Glow bijector, in preparation for converting Chain and Blockwise to CompositeTensor.

PiperOrigin-RevId: 375154471

Author: emilyfertig

tensorflow_probability/python/bijectors/composition.py

@@ -18,7 +18,6 @@
 from __future__ import division
 from __future__ import print_function
 
-import abc
 import collections
 import functools
 import sys
@@ -115,7 +114,11 @@ def __init__(self,
     """Instantiates a Composition of bijectors.
 
     Args:
-      bijectors: A nest-compatible structure of bijector instances.
+      bijectors: A nest-compatible structure of bijector instances or a
+        `Composition` bijector. If `bijectors` is a nested structure, then
+        `_walk_forward` and `_walk_inverse` must be implemented. If `bijectors`
+        is a `Composition` bijector, `_walk_forward` and `_walk_inverse` call
+        its corresponding methods.
       name: Name of this bijector.
       parameters: Dictionary of parameters used to initialize this bijector.
         These must be the exact values passed to `__init__`.
@@ -195,6 +198,8 @@ def _parameter_properties(cls, dtype):
 
   @property
   def bijectors(self):
+    if isinstance(self._bijectors, Composition):
+      return self._bijectors.bijectors  # pylint: disable=protected-access
     return self._bijectors
 
   @property
@@ -346,9 +351,6 @@ def transform_wrapper(bij, packed_ys, **nested):
         transform_wrapper, packed_args, **kwargs)
     return unpack_structs_like(self.forward_min_event_ndims, packed_result)
 
-  ### Abstract Methods
-
-  @abc.abstractmethod
   def _walk_forward(self, step_fn, argument, **kwargs):
     """Subclass stub for forward-mode traversals.
 
@@ -363,11 +365,17 @@ def _walk_forward(self, step_fn, argument, **kwargs):
         `_call_walk_forward` instead.
       argument: A (structure of) Tensor matching `self.forward_min_event_ndims`.
       **kwargs: Keyword arguments to be forwarded to nested bijectors.
+    Returns:
+      bijectors_forward: The value returned by `self._bijectors._walk_forward`
+        if `self._bijectors` is a `Composition` bijector.
+    Raises:
+      NotImplementedError, if `self._bijectors` is a nested structure.
     """
+    if isinstance(self._bijectors, Composition):
+      return self._bijectors._walk_forward(step_fn, argument, **kwargs)  # pylint: disable=protected-access
     raise NotImplementedError('{}._walk_forward is not implemented'.format(
         type(self).__name__))
 
-  @abc.abstractmethod
   def _walk_inverse(self, step_fn, argument, **kwargs):
     """Subclass stub for inverse-mode traversals.
 
@@ -382,7 +390,14 @@ def _walk_inverse(self, step_fn, argument, **kwargs):
         `_call_walk_inverse` instead.
       argument: A (structure of) Tensor matching `self.inverse_min_event_ndims`.
       **kwargs: Keyword arguments to be forwarded to nested bijectors.
+    Returns:
+      bijectors_inverse: The value returned by `self._bijectors._walk_inverse`
+        if `self._bijectors` is a `Composition` bijector.
+    Raises:
+      NotImplementedError, if `self._bijectors` is a nested structure.
     """
+    if isinstance(self._bijectors, Composition):
+      return self._bijectors._walk_inverse(step_fn, argument, **kwargs)  # pylint: disable=protected-access
     raise NotImplementedError('{}._walk_inverse is not implemented'.format(
         type(self).__name__))
 

tensorflow_probability/python/bijectors/glow.py

@@ -25,6 +25,7 @@
 from tensorflow_probability.python.bijectors import bijector
 from tensorflow_probability.python.bijectors import blockwise
 from tensorflow_probability.python.bijectors import chain
+from tensorflow_probability.python.bijectors import composition
 from tensorflow_probability.python.bijectors import exp
 from tensorflow_probability.python.bijectors import identity
 from tensorflow_probability.python.bijectors import invert
@@ -50,7 +51,7 @@
 ]
 
 
-class Glow(chain.Chain):
+class Glow(composition.Composition):
   r"""Implements the Glow Bijector from Kingma & Dhariwal (2018)[1].
 
   Overview: `Glow` is a chain of bijectors which transforms a rank-1 tensor
@@ -319,6 +320,7 @@ def __init__(self,
       name: Python `str`, name given to ops managed by this object.
         Default value: `'glow'`.
     """
+    parameters = dict(locals())
     # Make sure that the input shape is fully defined.
     if not tensorshape_util.is_fully_defined(output_shape):
       raise ValueError('Shape must be fully defined.')
@@ -441,10 +443,13 @@ def __init__(self,
           ]))
 
     glow_chain = glow_chain[::-1]
-    # To finish off, we initialize the bijector with the chain we've built
-    # This way, the rest of the model attributes are taken care of for us.
+    # To finish off, we build a bijector that chains the components together
+    # sequentially.
     super(Glow, self).__init__(
-        bijectors=glow_chain, validate_args=validate_args, name=name)
+        bijectors=chain.Chain(glow_chain, validate_args=validate_args),
+        validate_args=validate_args,
+        parameters=parameters,
+        name=name)
 
   @classmethod
   def _parameter_properties(cls, dtype):
@@ -491,7 +496,7 @@ def blockwise_splits(self):
     return self._blockwise_splits
 
 
-class ExitBijector(blockwise.Blockwise):
+class ExitBijector(composition.Composition):
   """The spatial coupling bijector used in Glow.
 
   This bijector consists of a blockwise bijector of a realNVP bijector. It is
@@ -524,7 +529,7 @@ def __init__(self,
         tensor with event shape `input_shape`, and returns a tensor with shape
         `input_shape`.
     """
-
+    parameters = dict(locals())
     nleave, ngrab, npass = blockwise_splits
 
     new_input_shape = input_shape[:-1]+(nleave,)
@@ -551,7 +556,10 @@ def __init__(self,
           bijector_fn=exit_bijector_fn)
 
     super(ExitBijector, self).__init__(
-        [shift_distribution, identity.Identity()], [nleave + ngrab, npass])
+        blockwise.Blockwise(
+            [shift_distribution, identity.Identity()], [nleave + ngrab, npass]),
+        parameters=parameters,
+        name='exit_bijector')
 
   @staticmethod
   def make_bijector_fn(layer, target_shape, scale_fn=tf.nn.sigmoid):
@@ -601,7 +609,7 @@ def bijector_fn(inputs, ignored_input):
     return bijector_fn
 
 
-class GlowBlock(chain.Chain):
+class GlowBlock(composition.Composition):
   """Single block for a glow model.
 
   This bijector contains `num_steps` steps of the flow, each consisting of an
@@ -613,7 +621,7 @@ class GlowBlock(chain.Chain):
 
   def __init__(self, input_shape, num_steps, coupling_bijector_fn,
                use_actnorm, seedstream):
-
+    parameters = dict(locals())
     rnvp_block = [identity.Identity()]
     this_nchan = input_shape[-1]
 
@@ -646,7 +654,8 @@ def __init__(self, input_shape, num_steps, coupling_bijector_fn,
 
     # Note that we reverse the list since Chain applies bijectors in reverse
     # order.
-    super(GlowBlock, self).__init__(rnvp_block[::-1])
+    super(GlowBlock, self).__init__(
+        chain.Chain(rnvp_block[::-1]), parameters=parameters, name='glow_block')
 
   @staticmethod
   def make_bijector_fn(layer, scale_fn=tf.nn.sigmoid):
@@ -809,7 +818,7 @@ def _init():
     return tf.cond(self._initialized, tf.no_op, _init)
 
 
-class Expand(chain.Chain):
+class Expand(composition.Composition):
   """A bijector to transform channels into spatial pixels."""
 
   def __init__(self, input_shape, block_size=2, validate_args=False, name=None):
@@ -827,8 +836,10 @@ def __init__(self, input_shape, block_size=2, validate_args=False, name=None):
             event_shape_in=[h, w, c],
             event_shape_out=[h, w, c // n**2, n, n]),
     ]
-    super(Expand, self).__init__(b, name=name or 'Expand',
-                                 parameters=parameters)
+    super(Expand, self).__init__(
+        bijectors=chain.Chain(b, validate_args=validate_args),
+        name=name or 'Expand',
+        parameters=parameters)
 
 
 class GlowDefaultNetwork(tfk.Sequential):

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 composition
 from tensorflow_probability.python.internal import test_util
 from tensorflow_probability.python.math.gradient import batch_jacobian
 
@@ -131,7 +132,7 @@ def testDataInit_inverse(self):
         x = tf.concat([x1, x2], axis=-1)
         nblocks += 1
 
-      elif isinstance(b, tfb.chain.Chain):
+      elif isinstance(b, composition.Composition):
         for bb in b.bijectors:
           x = self.evaluate(bb.inverse(x))
           if isinstance(bb, tfb.glow.ActivationNormalization):
@@ -172,7 +173,7 @@ def testDataInit_forward(self):
         y = tf.concat([y1, y2], axis=-1)
         nblocks += 1
 
-      elif isinstance(b, tfb.chain.Chain):
+      elif isinstance(b, composition.Composition):
         for bb in reversed(b.bijectors):
           y = self.evaluate(bb.forward(y))
           if isinstance(bb, tfb.glow.ActivationNormalization):