Hi!

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

Thanks!

Hi!

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

Thanks!

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

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 https://github.com/tensorflow/probability/blob/main/discussion/neutra/neutra_kernel.py) in TensorFlow:

```
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
```

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/affine_autoregressive.py`

. 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