To train:

python train.py [-h] (--basic | --nf | --iaf | --hf | --liaf) [--flow FLOW]

# E.g.
python train.py --basic
python train.py --nf --flow 10

• All the models (IAF, NF, VAE etc) are in models.py.

• All the neural net functions are in neural_networks.py.

• All the loss functions are in loss.py

• The train function accepts an encoder and a decoder. This is where you get to specify the type of encoder/decoder

• An encoder takes in (x, e) and spits out z.

• A decoder takes in z and spits out x

• To implement a new encoder F with hyperparameters W, we define a hidden function _F_encoder(x, e, W) and then set F_encoder(W) = lambda x, e: _F_encoder(x, e, W). Usually W comprises a neural network, flow lengths and so on. The method for decoders is identical. This allows us to define completely generic encoders and decoders with arbitrary structures and hyperparameters.