class_docstr: The spatial coupling bijector used in Glow.
  
  This bijector consists of a blockwise bijector of a realNVP bijector. It is
  where Glow adds a fork between points that are split off and passed to the
  base distribution, and points that are passed onward through more Glow blocks.
  
  For this bijector, we include spatial coupling between the part being forked
  off, and the part being passed onward. This induces a hierarchical spatial
  dependence on samples, and results in images which look better.
class_docstr_tok: ["The","spatial","coupling","bijector","used","in","Glow",".","This","bijector","consists","of","a","blockwise","bijector","of","a","realNVP","bijector",".","It","is","where","Glow","adds","a","fork","between","points","that","are","split","off","and","passed","to","the","base","distribution",",","and","points","that","are","passed","onward","through","more","Glow","blocks",".","For","this","bijector",",","we","include","spatial","coupling","between","the","part","being","forked","off",",","and","the","part","being","passed","onward",".","This","induces","a","hierarchical","spatial","dependence","on","samples",",","and","results","in","images","which","look","better","."]
docstr: Creates the exit bijector.
  
  Args:
    input_shape: A list specifying the input shape to the exit bijector.
      Used in constructing the network.
    blockwise_splits: A list of integers specifying the number of channels
      exiting the model, as well as those being left in the model, and those
      bypassing the exit bijector altogether.
    coupling_bijector_fn: A function which takes the argument `input_shape`
      and returns a callable neural network (e.g. a keras Sequential). The
      network should either return a tensor with the same event shape as
      `input_shape` (this will employ additive coupling), a tensor with the
      same height and width as `input_shape` but twice the number of channels
      (this will employ affine coupling), or a bijector which takes in a
      tensor with event shape `input_shape`, and returns a tensor with shape
      `input_shape`.
code: def __init__(
      self, input_shape, blockwise_splits, coupling_bijector_fn=None
  ):
      """Creates the exit bijector.
  
      Args:
        input_shape: A list specifying the input shape to the exit bijector.
          Used in constructing the network.
        blockwise_splits: A list of integers specifying the number of channels
          exiting the model, as well as those being left in the model, and those
          bypassing the exit bijector altogether.
        coupling_bijector_fn: A function which takes the argument `input_shape`
          and returns a callable neural network (e.g. a keras Sequential). The
          network should either return a tensor with the same event shape as
          `input_shape` (this will employ additive coupling), a tensor with the
          same height and width as `input_shape` but twice the number of channels
          (this will employ affine coupling), or a bijector which takes in a
          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,)
      target_output_shape = input_shape[:-1] + (ngrab,)
  
      # if nleave or ngrab == 0, then just use an identity for everything.
      if nleave == 0 or ngrab == 0:
          exit_layer = None
          exit_bijector_fn = None
  
          self.exit_layer = exit_layer
          shift_distribution = identity.Identity()
  
      else:
          exit_layer = coupling_bijector_fn(
              new_input_shape, output_chan=ngrab
          )
          exit_bijector_fn = self.make_bijector_fn(
              exit_layer,
              target_shape=target_output_shape,
              scale_fn=tf.exp,
          )
          self.exit_layer = exit_layer  # For variable tracking.
          shift_distribution = real_nvp.RealNVP(
              num_masked=nleave, bijector_fn=exit_bijector_fn
          )
  
      super(ExitBijector, self).__init__(
          blockwise.Blockwise(
              [shift_distribution, identity.Identity()],
              [nleave + ngrab, npass],
          ),
          parameters=parameters,
          name="exit_bijector",
      )
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  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
  from tensorflow_probability.python.bijectors import real_nvp
  from tensorflow_probability.python.bijectors import reshape
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import scale_matvec_lu
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.bijectors import transpose
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.util.deferred_tensor import TransformedVariable
  from tensorflow_probability.python.util.seed_stream import SeedStream
label_desc: Use image node_id 1 to create a new ExitBijector object from inherited base classes: composition with example: obj = ExitBijector(input_shape, blockwise_splits, coupling_bijector_fn)