nflows-extended: An extension of nflows

The original nflows package is a comprehensive collection of normalizing flows using PyTorch. Our nflows-extended, or enflows, is as the name suggests an extension. The main focus lies in implementing more flow layers from the literature in one consistent framework. In particular, we support conditional transformations based on hypernetworks. In the original package, conditional networks were restricted to using a conditional base distribution. In enflows, nearly every layer can be conditional :).

The bijective layers we additionally provide includes but are not limited to Planar Flows, invertible ResNets/DenseNets, a variant of neural autoregressive flows, and a basic support of continuous normalizing flows (and FFJORD) based torchdiffeq package.

Setting up the Environment.

The .yml file for the environment is given in env/conda_env.yml, and can be created from the base directory via:

(base) $  conda env create --file env/conda_env.yml
(base) $  conda activate environment.yml
# This code should then work:
(enflows) $  python examples/conditional_toy_2d.py

Ideally, after a successfull install you should be able to run and pass the unit tests with:

(enflows) /lagrangian_flow_net$ pytest

Usage

As the core is based on nflows, its usage is similar. To define a flow:

from enflows import transforms, distributions, flows

# Define an invertible transformation.
transform = transforms.CompositeTransform([
    transforms.MaskedAffineAutoregressiveTransform(features=2, hidden_features=4),
    transforms.RandomPermutation(features=2)
])

# Define a base distribution.
base_distribution = distributions.StandardNormal(shape=[2])


# Combine into a flow.
flow = flows.Flow(transform=transform, distribution=base_distribution)

To evaluate log probabilities of inputs:

log_prob = flow.log_prob(inputs)

To sample from the flow:

samples = flow.sample(num_samples)

Additional examples of the workflow are provided in examples folder.

Changes and added features compared to nflows

The core logic of the code for LFlows (i.e. the nflows/ directory) is based on the nflows package. Added Layers / Flwos:

• (Conditional)Sum-of-Sigmoid Layers
• Cholesky Outer Product for flows on symmetric positive definite matrices
• Lipschitz Constrained invertible DenseNets In particular, we provide three ways to condition these of these transformations without affecting the invertibility.
• Transformations for which the inverse is only known to exist, but not available:
  • (Conditional) Planar Flow
  • (Conditional) Sylvester Flow
• Conditional Versions of existing non-conditional transformations from nflows. Can be found for imports at nflows.transforms.conditional.*:
  • LU Transform
  • Orthogonal Transforms based on parameterized Householder projections
  • SVD based on the Orthogonal transforms
  • Shift Transform
• Conditional Versions of existing auto-regressive Variations, i.e. getting rid of the autoregressive parts.
  • ConditionalPiecewiseRationalQuadraticTransform
  • ConditionalUMNNTransform