experimental/neutra/neutra_kernel.py
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 = tfd.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.AutoregressiveLayer(
        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
def make_iaf_stack(total_event_size,
                   num_hidden_layers=2,
                   seed=None,
                   dtype=tf.float32):
  seed = tfd.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.AutoregressiveLayer(
        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
def make_iaf_stack ( total_event_size , num_hidden_layers = 2 , seed = None , dtype = tf . float32 ) : seed = tfd . 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 . AutoregressiveLayer ( 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
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.
Creates an stacked IAF bijector .