code: def make_iaf_stack(
      total_event_size, num_hidden_layers=2, seed=None, dtype=tf.float32
  ):
      """Creates an stacked IAF bijector.
  
      This bijector operates on vector-valued events.
  
      Args:
        total_event_size: Number of dimensions to operate over.
        num_hidden_layers: How many hidden layers to use in each IAF.
        seed: Random seed for the initializers.
        dtype: DType for the variables.
  
      Returns:
        bijector: The created bijector.
      """
  
      seed = tfp.util.SeedStream(seed, "make_iaf_stack")
  
      def make_iaf():
          """Create an IAF."""
          initializer = tf.compat.v2.keras.initializers.VarianceScaling(
              2 * 0.01, seed=seed() % (2**31 - 1)
          )
  
          made = tfb.AutoregressiveNetwork(
              params=2,
              event_shape=[total_event_size],
              hidden_units=[total_event_size] * num_hidden_layers,
              activation=tf.nn.elu,
              kernel_initializer=initializer,
              dtype=dtype,
          )
  
          def shift_and_scale(x):
              # TODO(siege): Something is losing the static shape.
              x.set_shape(
                  x.shape.merge_with(
                      [None] * (x.shape.ndims - 1) + [total_event_size]
                  )
              )
              return tf.unstack(made(x), num=2, axis=-1)
  
          return tfb.Invert(
              tfb.MaskedAutoregressiveFlow(shift_and_scale)
          )
  
      def make_swap():
          """Create an swap."""
          permutation = list(reversed(range(total_event_size)))
          return tfb.Permute(permutation)
  
      bijector = make_iaf()
      bijector = make_swap()(bijector)
      bijector = make_iaf()(bijector)
      bijector = make_swap()(bijector)
      bijector = make_iaf()(bijector)
      bijector = make_swap()(bijector)
  
      return bijector
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow_probability
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import unnest
label_desc: Use image node_id 1 for calling a global function with example usage: make_iaf_stack(total_event_size, num_hidden_layers, seed, dtype) and returns: bijector, tfb, tfb, tf