Adding nonlinearity to AutoIAFNormal guide

Is there a way to add non-linearity to AutoIAFNormal such as ELU?


hello can you provide (significantly) more detail about what you mean?

Hi @martinjankowiak

Sorry about that, I would like to add ELU activation between IAF flows as done in the paper of Hoffman et al., 2019 ([1903.03704] NeuTra-lizing Bad Geometry in Hamiltonian Monte Carlo Using Neural Transport). I looked in the source codes and couldn’t find where I can add it.
And maybe it is also a good place to ask if doing
guide = AutoIAFNormal(model, hidden_dim, num_transforms)
is comparable to the following code (I took from in TensorFlow:

def make_iaf_stack(total_event_size,
  """Creates an stacked IAF bijector.

  This bijector operates on vector-valued events.

    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.

    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(
        hidden_units=[total_event_size] * num_hidden_layers,

    def shift_and_scale(x):
      # TODO(siege): Something is losing the static 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

Basically, I am trying to implement the TensorFlow code above in Pyro.

assuming your goal is to do variational inference i think you’ll want to use AutoNormalizingFlow and pass a custom transform to it. in particular take a look at AffineAutoregressive in distributions/transforms/ the way this API was setup there’s some amount of flexibility but not infinite flexibility so i don’t think you’ll be able to get what you want with e.g. 3 lines of code.

1 Like