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","."]
code: def make_bijector_fn(layer, target_shape, scale_fn=tf.nn.sigmoid):
      def bijector_fn(inputs, ignored_input):
          """Decorated function to get the RealNVP bijector."""
          # Build this so we can handle a user passing a NN that returns a tensor
          # OR an NN that returns a bijector
          possible_output = layer(inputs)
  
          # We need to produce a bijector, but we do not know if the layer has done
          # so. We are setting this up to handle 2 possibilities:
          # 1) The layer outputs a bijector --> all is good
          # 2) The layer outputs a tensor --> we need to turn it into a bijector.
          if isinstance(possible_output, bijector.Bijector):
              output = possible_output
          elif isinstance(possible_output, tf.Tensor):
              input_shape = inputs.get_shape().as_list()
              output_shape = possible_output.get_shape().as_list()
              assert input_shape[:-1] == output_shape[:-1]
              c = input_shape[-1]
  
              # For layers which output a tensor, we have two possibilities:
              # 1) There are twice as many output channels as the target --> the
              #    coupling is affine, meaning there is a scale followed by a shift.
              # 2) The number of output channels equals the target --> the
              #    coupling is additive, meaning there is just a shift
              if target_shape[-1] == output_shape[-1] // 2:
                  this_scale = scale.Scale(
                      scale_fn(possible_output[..., :c] + 2.0)
                  )
                  this_shift = shift.Shift(possible_output[..., c:])
                  output = this_shift(this_scale)
              elif target_shape[-1] == output_shape[-1]:
                  output = shift.Shift(possible_output[..., :c])
              else:
                  raise ValueError(
                      "Shape inconsistent with input. Expected shape"
                      "{0} or {1} but tensor was shape {2}".format(
                          input_shape,
                          tf.concat(
                              [input_shape[:-1], [2 * input_shape[-1]]],
                              0,
                          ),
                          output_shape,
                      )
                  )
          else:
              raise ValueError(
                  "Expected a bijector or a tensor, but instead got"
                  "{}".format(possible_output.__class__)
              )
          return output
  
      return bijector_fn
  
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 2 for calling the ExitBijector obj's underlying member method code with example usage: obj.make_bijector_fn(layer, target_shape, scale_fn) and returns: bijector_fn, output