class_docstr: Implements the Glow Bijector from Kingma & Dhariwal (2018)[1].
  
  Overview: `Glow` is a chain of bijectors which transforms a rank-1 tensor
  (vector) into a rank-3 tensor (e.g. an RGB image). `Glow` does this by
  chaining together an alternating series of "Blocks," "Squeezes," and "Exits"
  which are each themselves special chains of other bijectors. The intended use
  of `Glow` is as part of a `tfp.distributions.TransformedDistribution`, in
  which the base distribution over the vector space is used to generate samples
  in the image space. In the paper, an Independent Normal distribution is used
  as the base distribution.
  
  A "Block" (implemented as the `GlowBlock` Bijector) performs much of the
  transformations which allow glow to produce sophisticated and complex mappings
  between the image space and the latent space and therefore achieve rich image
  generation performance. A Block is composed of `num_steps_per_block` steps,
  which are each implemented as a `Chain` containing an
  `ActivationNormalization` (ActNorm) bijector, followed by an (invertible)
  `OneByOneConv` bijector, and finally a coupling bijector. The coupling
  bijector is an instance of a `RealNVP` bijector, and uses the
  `coupling_bijector_fn` function to instantiate the coupling bijector function
  which is given to the `RealNVP`. This function returns a bijector which
  defines the coupling (e.g. `Shift(Scale)` for affine coupling or `Shift` for
  additive coupling).
  
  A "Squeeze" converts spatial features into channel features. It is
  implemented using the `Expand` bijector. The difference in names is
  due to the fact that the `forward` function from glow is meant to ultimately
  correspond to sampling from a `tfp.util.TransformedDistribution` object,
  which would use `Expand` (Squeeze is just Invert(Expand)). The `Expand`
  bijector takes a tensor with shape `[H, W, C]` and returns a tensor with shape
  `[2H, 2W, C / 4]`, such that each 2x2x1 spatial tile in the output is composed
  from a single 1x1x4 tile in the input tensor, as depicted in the figure below.
  
                           Forward pass (Expand)
                          ______        __________
                          \     \       \    \    \
                          \\     \ ----> \  1 \  2 \
                          \\\__1__\       \____\____\
                          \\\__2__\        \    \    \
                           \\__3__\  <----  \  3 \  4 \
                            \__4__\          \____\____\
                               Inverse pass (Squeeze)
  
  This is implemented using a chain of `Reshape` -> `Transpose` -> `Reshape`
  bijectors. Note that on an inverse pass through the bijector, each Squeeze
  will cause the width/height of the image to decrease by a factor of 2.
  Therefore, the input image must be evenly divisible by 2 at least
  `num_glow_blocks` times, since it will pass through a Squeeze step that many
  times.
  
  An "Exit" is simply a junction at which some of the tensor "exits" from the
  glow bijector and therefore avoids any further alteration. Each exit is
  implemented as a `Blockwise` bijector, where some channels are given to the
  rest of the glow model, and the rest are given to a bypass implemented using
  the `Identity` bijector. The fraction of channels to be removed at each exit
  is determined by the `grab_after_block` arg, indicates the fraction of
  _remaining_ channels which join the identity bypass. The fraction is
  converted to an integer number of channels by multiplying by the remaining
  number of channels and rounding.
  
  Additionally, at each exit, glow couples the tensor exiting the highway to
  the tensor continuing onward. This makes small scale features in the image
  dependent on larger scale features, since the larger scale features dictate
  the mean and scale of the distribution over the smaller scale features.
  This coupling is done similarly to the Coupling bijector in each step of the
  flow (i.e. using a RealNVP bijector). However for the exit bijector, the
  coupling is instantiated using `exit_bijector_fn` rather than coupling
  bijector fn, allowing for different behaviors between standard coupling and
  exit coupling. Also note that because the exit utilizes a coupling bijector,
  there are two special cases (all channels exiting and no channels exiting).
  
  The full Glow bijector consists of `num_glow_blocks` Blocks each of which
  contains `num_steps_per_block` steps. Each step implements a coupling using
  `bijector_coupling_fn`. Between blocks, glow converts between spatial pixels
  and channels using the Expand Bijector, and splits channels out of the
  bijector using the Exit Bijector. The channels which have exited continue
  onward through Identity bijectors and those which have not exited are given
  to the next block. After passing through all Blocks, the tensor is reshaped
  to a rank-1 tensor with the same number of elements. This is where the
  distribution will be defined.
  
  A schematic diagram of Glow is shown below. The `forward` function of the
  bijector starts from the bottom and goes upward, while the `inverse` function
  starts from the top and proceeds downward.
  
  ```None
  ==============================================================================
                           Glow Schematic Diagram
  
  Input Image     ########################   shape = [H, W, C]
  
                  \                      /<- Expand Bijector turns spatial
                   \                    /    dimensions into channels.
                  _
                 |  XXXXXXXXXXXXXXXXXXXX
                 |  XXXXXXXXXXXXXXXXXXXX
                 |  XXXXXXXXXXXXXXXXXXXX     A single step of the flow consists
   Glow Block  - |  XXXXXXXXXXXXXXXXXXXX  <- of ActNorm -> 1x1Conv -> Coupling.
                 |  XXXXXXXXXXXXXXXXXXXX     there are num_steps_per_block
                 |  XXXXXXXXXXXXXXXXXXXX     steps of the flow in each block.
                 |_ XXXXXXXXXXXXXXXXXXXX
  
                    \                  / <-- Expand bijectors follow each glow
                     \                /      block
  
                      XXXXXXXX\\\\\\\\   <-- Exit Bijector removes channels
                  _                    _     from additional alteration.
                 |    XXXXXXXX !  |  !
                 |    XXXXXXXX !  |  !
                 |    XXXXXXXX !  |  !       After exiting, channels are passed
   Glow Block  - |    XXXXXXXX !  |  !  <--- downward using the Blockwise and
                 |    XXXXXXXX !  |  !       Identify bijectors.
                 |    XXXXXXXX !  |  !
                 |_   XXXXXXXX !  |  !
  
                      \              / <---- Expand Bijector
                       \            /
  
                        XXX\\\    | !  <---- Exit Bijector
                  _
                 |      XXX ! |   | !
                 |      XXX ! |   | !
                 |      XXX ! |   | !
   Glow Block  - |      XXX ! |   | !
                 |      XXX ! |   | !
                 |      XXX ! |   | !
                 |_     XXX ! |   | !
  
                        XX\ ! |   | ! <----- (Optional) Exit Bijector
  
                         |    |   |
                         v    v   v
  Output Distribution    ##########          shape = [H * W * C]
                                                     _________________________
                                                    |         Legend          |
                                                    | XX  = Step of flow      |
                                                    | X\  = Exit bijector     |
                                                    | \/  = Expand bijector   |
                                                    | !|! = Identity bijector |
                                                    |                         |
                                                    | up  = Forward pass      |
                                                    | dn  = Inverse pass      |
                                                    |_________________________|
  
  ==============================================================================
  ```
  The default configuration for glow is meant to replicate the architecture in
  [1] for generating images from CIFAR-10.
  
  Example usage:
  ```python
  
  from functools import reduce
  from operator import mul
  import tensorflow as tf
  import tensorflow_datasets as tfds
  import tensorflow_probability as tfp
  tfb = tfp.bijectors
  tfd = tfp.distributions
  
  data, info = tfds.load('cifar10', with_info=True)
  train_data, test_data = data['train'], data['test']
  
  preprocess = lambda x: tf.cast(x['image'], tf.float32)
  train_data = train_data.batch(4).map(preprocess)
  test_data = test_data.batch(4).map(preprocess)
  
  x = next(iter(train_data))
  
  glow = tfb.Glow(output_shape=info.features['image'].shape,
                  coupling_bijector_fn=tfb.GlowDefaultNetwork,
                  exit_bijector_fn=tfb.GlowDefaultExitNetwork)
  
  z_shape = glow.inverse_event_shape(info.features['image'].shape)
  
  pz = tfd.Sample(tfd.Normal(0., 1.), z_shape)
  
  # Calling glow on distribution p(z) creates our glow distribution over images.
  px = glow(pz)
  
  # Take samples from the distribution to get images from your dataset
  images = px.sample(4)
  
  # Map images to positions in the distribution
  z = glow.inverse(x)
  
  # Get the z's corresponding to each spatial scale. To do this, we have to
  # find out how many zs are passed through blockwise at each stage that were
  # not passed at the previous stage. This is encoded in the second element of
  # each list of blockwise splits. However because the bijector iteratively
  # converts spatial pixels to channels, we also need to multiply the size of
  # that second element by the number of spatial-to-channel conversions that the
  # tensor receives after exiting (including after any alteration).
  ztake = [bs[1] * 4**(i+2) for i, bs in enumerate(glow.blockwise_splits)]
  total_z_taken = sum(ztake)
  split_sizes = [z_shape.as_list()[0]-total_z_taken] + ztake
  zsplits = tf.split(z, num_or_size_splits=split_sizes, axis=-1)
  ```
  
  #### References:
  
  [1]: Diederik P Kingma, Prafulla Dhariwal, Glow: Generative Flow
       with Invertible 1x1 Convolutions. In _Neural Information
       Processing Systems_, 2018. https://arxiv.org/abs/1807.03039
  
  [2]: Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density Estimation
       using Real NVP. In _International Conference on Learning
       Representations_, 2017. https://arxiv.org/abs/1605.08803
class_docstr_tok: ["Implements","the","Glow","Bijector","from","Kingma","&","Dhariwal","(","2018",")","[","1","]",".","Overview",":","`","Glow","`","is","a","chain","of","bijectors","which","transforms","a","rank-1","tensor","(","vector",")","into","a","rank-3","tensor","(","e.g",".","an","RGB","image",")",".","`","Glow","`","does","this","by","chaining","together","an","alternating","series","of","``","Blocks",",","''","``","Squeezes",",","''","and","``","Exits","''","which","are","each","themselves","special","chains","of","other","bijectors",".","The","intended","use","of","`","Glow","`","is","as","part","of","a","`","tfp.distributions.TransformedDistribution","`",",","in","which","the","base","distribution","over","the","vector","space","is","used","to","generate","samples","in","the","image","space",".","In","the","paper",",","an","Independent","Normal","distribution","is","used","as","the","base","distribution",".","A","``","Block","''","(","implemented","as","the","`","GlowBlock","`","Bijector",")","performs","much","of","the","transformations","which","allow","glow","to","produce","sophisticated","and","complex","mappings","between","the","image","space","and","the","latent","space","and","therefore","achieve","rich","image","generation","performance",".","A","Block","is","composed","of","`","num_steps_per_block","`","steps",",","which","are","each","implemented","as","a","`","Chain","`","containing","an","`","ActivationNormalization","`","(","ActNorm",")","bijector",",","followed","by","an","(","invertible",")","`","OneByOneConv","`","bijector",",","and","finally","a","coupling","bijector",".","The","coupling","bijector","is","an","instance","of","a","`","RealNVP","`","bijector",",","and","uses","the","`","coupling_bijector_fn","`","function","to","instantiate","the","coupling","bijector","function","which","is","given","to","the","`","RealNVP","`",".","This","function","returns","a","bijector","which","defines","the","coupling","(","e.g",".","`","Shift","(","Scale",")","`","for","affine","coupling","or","`","Shift","`","for","additive","coupling",")",".","A","``","Squeeze","''","converts","spatial","features","into","channel","features",".","It","is","implemented","using","the","`","Expand","`","bijector",".","The","difference","in","names","is","due","to","the","fact","that","the","`","forward","`","function","from","glow","is","meant","to","ultimately","correspond","to","sampling","from","a","`","tfp.util.TransformedDistribution","`","object",",","which","would","use","`","Expand","`","(","Squeeze","is","just","Invert","(","Expand",")",")",".","The","`","Expand","`","bijector","takes","a","tensor","with","shape","`","[","H",",","W",",","C","]","`","and","returns","a","tensor","with","shape","`","[","2H",",","2W",",","C","\/","4","]","`",",","such","that","each","2x2x1","spatial","tile","in","the","output","is","composed","from","a","single","1x1x4","tile","in","the","input","tensor",",","as","depicted","in","the","figure","below",".","Forward","pass","(","Expand",")","______","__________","\\","\\","\\","\\","\\","\\\\","\\","--","--",">","\\","1","\\","2","\\","\\\\\\__1__\\","\\____\\____\\","\\\\\\__2__\\","\\","\\","\\","\\\\__3__\\","<","--","--","\\","3","\\","4","\\","\\__4__\\","\\____\\____\\","Inverse","pass","(","Squeeze",")","This","is","implemented","using","a","chain","of","`","Reshape","`","-",">","`","Transpose","`","-",">","`","Reshape","`","bijectors",".","Note","that","on","an","inverse","pass","through","the","bijector",",","each","Squeeze","will","cause","the","width\/height","of","the","image","to","decrease","by","a","factor","of","2",".","Therefore",",","the","input","image","must","be","evenly","divisible","by","2","at","least","`","num_glow_blocks","`","times",",","since","it","will","pass","through","a","Squeeze","step","that","many","times",".","An","``","Exit","''","is","simply","a","junction","at","which","some","of","the","tensor","``","exits","''","from","the","glow","bijector","and","therefore","avoids","any","further","alteration",".","Each","exit","is","implemented","as","a","`","Blockwise","`","bijector",",","where","some","channels","are","given","to","the","rest","of","the","glow","model",",","and","the","rest","are","given","to","a","bypass","implemented","using","the","`","Identity","`","bijector",".","The","fraction","of","channels","to","be","removed","at","each","exit","is","determined","by","the","`","grab_after_block","`","arg",",","indicates","the","fraction","of","_remaining_","channels","which","join","the","identity","bypass",".","The","fraction","is","converted","to","an","integer","number","of","channels","by","multiplying","by","the","remaining","number","of","channels","and","rounding",".","Additionally",",","at","each","exit",",","glow","couples","the","tensor","exiting","the","highway","to","the","tensor","continuing","onward",".","This","makes","small","scale","features","in","the","image","dependent","on","larger","scale","features",",","since","the","larger","scale","features","dictate","the","mean","and","scale","of","the","distribution","over","the","smaller","scale","features",".","This","coupling","is","done","similarly","to","the","Coupling","bijector","in","each","step","of","the","flow","(","i.e",".","using","a","RealNVP","bijector",")",".","However","for","the","exit","bijector",",","the","coupling","is","instantiated","using","`","exit_bijector_fn","`","rather","than","coupling","bijector","fn",",","allowing","for","different","behaviors","between","standard","coupling","and","exit","coupling",".","Also","note","that","because","the","exit","utilizes","a","coupling","bijector",",","there","are","two","special","cases","(","all","channels","exiting","and","no","channels","exiting",")",".","The","full","Glow","bijector","consists","of","`","num_glow_blocks","`","Blocks","each","of","which","contains","`","num_steps_per_block","`","steps",".","Each","step","implements","a","coupling","using","`","bijector_coupling_fn","`",".","Between","blocks",",","glow","converts","between","spatial","pixels","and","channels","using","the","Expand","Bijector",",","and","splits","channels","out","of","the","bijector","using","the","Exit","Bijector",".","The","channels","which","have","exited","continue","onward","through","Identity","bijectors","and","those","which","have","not","exited","are","given","to","the","next","block",".","After","passing","through","all","Blocks",",","the","tensor","is","reshaped","to","a","rank-1","tensor","with","the","same","number","of","elements",".","This","is","where","the","distribution","will","be","defined",".","A","schematic","diagram","of","Glow","is","shown","below",".","The","`","forward","`","function","of","the","bijector","starts","from","the","bottom","and","goes","upward",",","while","the","`","inverse","`","function","starts","from","the","top","and","proceeds","downward",".","``","`","None","==============================================================================","Glow","Schematic","Diagram","Input","Image","#","#","#","#","#","#","#","#","#","#","#","#","#","#","#","#","#","#","#","#","#","#","#","#","shape","=","[","H",",","W",",","C","]","\\","\/","<","-","Expand","Bijector","turns","spatial","\\","\/","dimensions","into","channels",".","_","|","XXXXXXXXXXXXXXXXXXXX","|","XXXXXXXXXXXXXXXXXXXX","|","XXXXXXXXXXXXXXXXXXXX","A","single","step","of","the","flow","consists","Glow","Block","-","|","XXXXXXXXXXXXXXXXXXXX","<","-","of","ActNorm","-",">","1x1Conv","-",">","Coupling",".","|","XXXXXXXXXXXXXXXXXXXX","there","are","num_steps_per_block","|","XXXXXXXXXXXXXXXXXXXX","steps","of","the","flow","in","each","block",".","|_","XXXXXXXXXXXXXXXXXXXX","\\","\/","<","--","Expand","bijectors","follow","each","glow","\\","\/","block","XXXXXXXX\\\\\\\\\\\\\\\\","<","--","Exit","Bijector","removes","channels","_","_","from","additional","alteration",".","|","XXXXXXXX","!","|","!","|","XXXXXXXX","!","|","!","|","XXXXXXXX","!","|","!","After","exiting",",","channels","are","passed","Glow","Block","-","|","XXXXXXXX","!","|","!","<","--","-","downward","using","the","Blockwise","and","|","XXXXXXXX","!","|","!","Identify","bijectors",".","|","XXXXXXXX","!","|","!","|_","XXXXXXXX","!","|","!","\\","\/","<","--","--","Expand","Bijector","\\","\/","XXX\\\\\\","|","!","<","--","--","Exit","Bijector","_","|","XXX","!","|","|","!","|","XXX","!","|","|","!","|","XXX","!","|","|","!","Glow","Block","-","|","XXX","!","|","|","!","|","XXX","!","|","|","!","|","XXX","!","|","|","!","|_","XXX","!","|","|","!","XX\\","!","|","|","!","<","--","--","-","(","Optional",")","Exit","Bijector","|","|","|","v","v","v","Output","Distribution","#","#","#","#","#","#","#","#","#","#","shape","=","[","H","*","W","*","C","]","_________________________","|","Legend","|","|","XX","=","Step","of","flow","|","|","X\\","=","Exit","bijector","|","|","\\\/","=","Expand","bijector","|","|","!","|","!","=","Identity","bijector","|","|","|","|","up","=","Forward","pass","|","|","dn","=","Inverse","pass","|","|_________________________|","==============================================================================","``","`","The","default","configuration","for","glow","is","meant","to","replicate","the","architecture","in","[","1","]","for","generating","images","from","CIFAR-10",".","Example","usage",":","``","`","python","from","functools","import","reduce","from","operator","import","mul","import","tensorflow","as","tf","import","tensorflow_datasets","as","tfds","import","tensorflow_probability","as","tfp","tfb","=","tfp.bijectors","tfd","=","tfp.distributions","data",",","info","=","tfds.load","(","'cifar10","'",",","with_info=True",")","train_data",",","test_data","=","data","[","'train","'","]",",","data","[","'test","'","]","preprocess","=","lambda","x",":","tf.cast","(","x","[","'image","'","]",",","tf.float32",")","train_data","=","train_data.batch","(","4",")",".map","(","preprocess",")","test_data","=","test_data.batch","(","4",")",".map","(","preprocess",")","x","=","next","(","iter","(","train_data",")",")","glow","=","tfb.Glow","(","output_shape=info.features","[","'image","'","]",".shape",",","coupling_bijector_fn=tfb.GlowDefaultNetwork",",","exit_bijector_fn=tfb.GlowDefaultExitNetwork",")","z_shape","=","glow.inverse_event_shape","(","info.features","[","'image","'","]",".shape",")","pz","=","tfd.Sample","(","tfd.Normal","(","0.",",","1",".",")",",","z_shape",")","#","Calling","glow","on","distribution","p","(","z",")","creates","our","glow","distribution","over","images",".","px","=","glow","(","pz",")","#","Take","samples","from","the","distribution","to","get","images","from","your","dataset","images","=","px.sample","(","4",")","#","Map","images","to","positions","in","the","distribution","z","=","glow.inverse","(","x",")","#","Get","the","z","'s","corresponding","to","each","spatial","scale",".","To","do","this",",","we","have","to","#","find","out","how","many","zs","are","passed","through","blockwise","at","each","stage","that","were","#","not","passed","at","the","previous","stage",".","This","is","encoded","in","the","second","element","of","#","each","list","of","blockwise","splits",".","However","because","the","bijector","iteratively","#","converts","spatial","pixels","to","channels",",","we","also","need","to","multiply","the","size","of","#","that","second","element","by","the","number","of","spatial-to-channel","conversions","that","the","#","tensor","receives","after","exiting","(","including","after","any","alteration",")",".","ztake","=","[","bs","[","1","]","*","4","*","*","(","i+2",")","for","i",",","bs","in","enumerate","(","glow.blockwise_splits",")","]","total_z_taken","=","sum","(","ztake",")","split_sizes","=","[","z_shape.as_list","(",")","[","0","]","-total_z_taken","]","+","ztake","zsplits","=","tf.split","(","z",",","num_or_size_splits=split_sizes",",","axis=-1",")","``","`","#","#","#","#","References",":","[","1","]",":","Diederik","P","Kingma",",","Prafulla","Dhariwal",",","Glow",":","Generative","Flow","with","Invertible","1x1","Convolutions",".","In","_Neural","Information","Processing","Systems_",",","2018.","https",":","\/\/arxiv.org\/abs\/1807.03039","[","2","]",":","Laurent","Dinh",",","Jascha","Sohl-Dickstein",",","and","Samy","Bengio",".","Density","Estimation","using","Real","NVP",".","In","_International","Conference","on","Learning","Representations_",",","2017.","https",":","\/\/arxiv.org\/abs\/1605.08803"]
docstr: Creates the Glow bijector.
  
  Args:
    output_shape: A list of integers, specifying the event shape of the
      output, of the bijectors forward pass (the image).  Specified as
      [H, W, C].
      Default Value: (32, 32, 3)
    num_glow_blocks: An integer, specifying how many downsampling levels to
      include in the model. This must divide equally into both H and W,
      otherwise the bijector would not be invertible.
      Default Value: 3
    num_steps_per_block: An integer specifying how many Affine Coupling and
      1x1 convolution layers to include at each level of the spatial
      hierarchy.
      Default Value: 32 (i.e. the value used in the original glow paper).
    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`.
    exit_bijector_fn: Similar to coupling_bijector_fn, exit_bijector_fn is
      a function which takes the argument `input_shape` and `output_chan`
      and returns a callable neural network. The neural network it returns
      should take a tensor of shape `input_shape` as the input, and return
      one of three options: A tensor with `output_chan` channels, a tensor
      with `2 * output_chan` channels, or a bijector. Additional details can
      be found in the documentation for ExitBijector.
    grab_after_block: A tuple of floats, specifying what fraction of the
      remaining channels to remove following each glow block. Glow will take
      the integer floor of this number multiplied by the remaining number of
      channels. The default is half at each spatial hierarchy.
      Default value: None (this will take out half of the channels after each
        block.
    use_actnorm: A bool deciding whether or not to use actnorm. Data-dependent
      initialization is used to initialize this layer.
      Default value: `False`
    seed: A seed to control randomness in the 1x1 convolution initialization.
      Default value: `None` (i.e., non-reproducible sampling).
    validate_args: Python `bool` indicating whether arguments should be
      checked for correctness.
      Default value: `False`
    name: Python `str`, name given to ops managed by this object.
      Default value: `'glow'`.
code: def __init__(
      self,
      output_shape=(32, 32, 3),
      num_glow_blocks=3,
      num_steps_per_block=32,
      coupling_bijector_fn=None,
      exit_bijector_fn=None,
      grab_after_block=None,
      use_actnorm=True,
      seed=None,
      validate_args=False,
      name="glow",
  ):
      """Creates the Glow bijector.
  
      Args:
        output_shape: A list of integers, specifying the event shape of the
          output, of the bijectors forward pass (the image).  Specified as
          [H, W, C].
          Default Value: (32, 32, 3)
        num_glow_blocks: An integer, specifying how many downsampling levels to
          include in the model. This must divide equally into both H and W,
          otherwise the bijector would not be invertible.
          Default Value: 3
        num_steps_per_block: An integer specifying how many Affine Coupling and
          1x1 convolution layers to include at each level of the spatial
          hierarchy.
          Default Value: 32 (i.e. the value used in the original glow paper).
        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`.
        exit_bijector_fn: Similar to coupling_bijector_fn, exit_bijector_fn is
          a function which takes the argument `input_shape` and `output_chan`
          and returns a callable neural network. The neural network it returns
          should take a tensor of shape `input_shape` as the input, and return
          one of three options: A tensor with `output_chan` channels, a tensor
          with `2 * output_chan` channels, or a bijector. Additional details can
          be found in the documentation for ExitBijector.
        grab_after_block: A tuple of floats, specifying what fraction of the
          remaining channels to remove following each glow block. Glow will take
          the integer floor of this number multiplied by the remaining number of
          channels. The default is half at each spatial hierarchy.
          Default value: None (this will take out half of the channels after each
            block.
        use_actnorm: A bool deciding whether or not to use actnorm. Data-dependent
          initialization is used to initialize this layer.
          Default value: `False`
        seed: A seed to control randomness in the 1x1 convolution initialization.
          Default value: `None` (i.e., non-reproducible sampling).
        validate_args: Python `bool` indicating whether arguments should be
          checked for correctness.
          Default value: `False`
        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.")
      if tensorshape_util.rank(output_shape) != 3:
          raise ValueError(
              "Shape ndims must be 3 for images.  Your shape is"
              "{}".format(tensorshape_util.rank(output_shape))
          )
  
      num_glow_blocks_ = tf.get_static_value(num_glow_blocks)
      if (
          num_glow_blocks_ is None
          or int(num_glow_blocks_) != num_glow_blocks_
          or num_glow_blocks_ < 1
      ):
          raise ValueError(
              "Argument `num_glow_blocks` must be a statically known"
              "positive `int` (saw: {}).".format(num_glow_blocks)
          )
      num_glow_blocks = int(num_glow_blocks_)
  
      output_shape = tensorshape_util.as_list(output_shape)
      h, w, c = output_shape
      n = num_glow_blocks
      nsteps = num_steps_per_block
  
      # Default Glow: Half of the channels are split off after each block,
      # and after the final block, no channels are split off.
      if grab_after_block is None:
          grab_after_block = tuple([0.5] * (n - 1) + [0.0])
  
      # Thing we know must be true: h and w are evenly divisible by 2, n times.
      # Otherwise, the squeeze bijector will not work.
      if w % 2**n != 0:
          raise ValueError(
              "Width must be divisible by 2 at least n times."
              "Saw: {} % {} != 0".format(w, 2**n)
          )
      if h % 2**n != 0:
          raise ValueError(
              "Height should be divisible by 2 at least n times."
          )
      if h // 2**n < 1:
          raise ValueError(
              "num_glow_blocks ({0}) is too large. The image height "
              "({1}) must be divisible by 2 no more than {2} "
              "times.".format(
                  num_glow_blocks, h, int(np.log(h) / np.log(2.0))
              )
          )
      if w // 2**n < 1:
          raise ValueError(
              "num_glow_blocks ({0}) is too large. The image width "
              "({1}) must be divisible by 2 no more than {2} "
              "times.".format(
                  num_glow_blocks, w, int(np.log(h) / np.log(2.0))
              )
          )
  
      # Other things we want to be true:
      # - The number of times we take must be equal to the number of glow blocks.
      if len(grab_after_block) != num_glow_blocks:
          raise ValueError(
              "Length of grab_after_block ({0}) must match the number"
              "of blocks ({1}).".format(
                  len(grab_after_block), num_glow_blocks
              )
          )
  
      self._blockwise_splits = self._get_blockwise_splits(
          output_shape, grab_after_block[::-1]
      )
  
      # Now check on the values of blockwise splits
      if any([bs[0] < 1 for bs in self._blockwise_splits]):
          first_offender = [
              bs[0] for bs in self._blockwise_splits
          ].index(True)
          raise ValueError(
              "At at least one exit, you are taking out all of your "
              "channels, and therefore have no inputs to later blocks."
              " Try setting grab_after_block to a lower value at index"
              "{}.".format(first_offender)
          )
  
      if any(np.isclose(gab, 0) for gab in grab_after_block):
          # Special case: if specifically exiting no channels, then the exit is
          # just an identity bijector.
          pass
      elif any([bs[1] < 1 for bs in self._blockwise_splits]):
          first_offender = [
              bs[1] for bs in self._blockwise_splits
          ].index(True)
          raise ValueError(
              "At least one of your layers has < 1 output channels. "
              "This means you set grab_at_block too small. "
              "Try setting grab_after_block to a larger value at index"
              "{}.".format(first_offender)
          )
  
      # Lets start to build our bijector. We assume that the distribution is 1
      # dimensional. First, lets reshape it to an image.
      glow_chain = [
          reshape.Reshape(
              event_shape_out=[h // 2**n, w // 2**n, c * 4**n],
              event_shape_in=[h * w * c],
          )
      ]
  
      seedstream = SeedStream(seed=seed, salt="random_beta")
  
      for i in range(n):
          # This is the shape of the current tensor
          current_shape = (
              h // 2**n * 2**i,
              w // 2**n * 2**i,
              c * 4 ** (i + 1),
          )
  
          # This is the shape of the input to both the glow block and exit bijector.
          this_nchan = sum(self._blockwise_splits[i][0:2])
          this_input_shape = (
              h // 2**n * 2**i,
              w // 2**n * 2**i,
              this_nchan,
          )
  
          glow_chain.append(
              invert.Invert(
                  ExitBijector(
                      current_shape,
                      self._blockwise_splits[i],
                      exit_bijector_fn,
                  )
              )
          )
  
          glow_block = GlowBlock(
              input_shape=this_input_shape,
              num_steps=nsteps,
              coupling_bijector_fn=coupling_bijector_fn,
              use_actnorm=use_actnorm,
              seedstream=seedstream,
          )
  
          if self._blockwise_splits[i][2] == 0:
              # All channels are passed to the RealNVP
              glow_chain.append(glow_block)
          else:
              # Some channels are passed around the block.
              # This is done with the Blockwise bijector.
              glow_chain.append(
                  blockwise.Blockwise(
                      [glow_block, identity.Identity()],
                      [
                          sum(self._blockwise_splits[i][0:2]),
                          self._blockwise_splits[i][2],
                      ],
                  )
              )
  
          # Finally, lets expand the channels into spatial features.
          glow_chain.append(
              Expand(
                  input_shape=[
                      h // 2**n * 2**i,
                      w // 2**n * 2**i,
                      c * 4**n // 4**i,
                  ]
              )
          )
  
      glow_chain = glow_chain[::-1]
      # To finish off, we build a bijector that chains the components together
      # sequentially.
      super(Glow, self).__init__(
          bijectors=chain.Chain(
              glow_chain, validate_args=validate_args
          ),
          validate_args=validate_args,
          parameters=parameters,
          name=name,
      )
  
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 Glow object from inherited base classes: composition with example: obj = Glow(output_shape, num_glow_blocks, num_steps_per_block, coupling_bijector_fn, exit_bijector_fn, grab_after_block, use_actnorm, seed, validate_args, name)