New Pytorch package - Jammy Flows

[Thorsten Glüsenkamp]

I would like to add the package jammy_flows to the collection (https://github.com/thoglu/jammy_flows). I developed it for normalizing-flow inference in the context of astro-particle physics (https://arxiv.org/abs/2008.05825), but I think it might be useful in general.

It stands for Joint Autoregressive Manifold (MY) flows and models a (conditional) PDF on a tensor product of manifolds. The sub-PDFs are connected autoregressively similar to Inverse Autoregressive Flows (IAF - Kingma et al. 2016), but compared to IAF generalize via 1) allowing for arbitrary (non-affine) coupling layers (every flow in the package is amortizable) and 2) because we have general couplings we can link flows on different manifolds (e.g. Euclidean and a sphere).

It is mostly designed for low-dimensional applications (maybe a few 10s of dimensions - although simple flows ĺike the affine flow should reasonably work at much higher dimensionality) and should be simple to set up.

For example, a 5-d PDF defined on a 3-dimensional Euclidean manifold and an autoregressively linked 2-sphere conditional PDF is defined like this (together this forms a joint distribution on the tensor product space $\mathcal{R}^3 \times \mathcal{S}^2)$:

import jammy_flows

pdf=jammy_flows.pdf("e3+s2", "gg+n")

The first argument defines the autoregressive manifold structure, the second argument the exact flow layer used for each manifold.
Each flow is abbreviated with a letter (see below) - for example "g" stand for a Gaussinization flow layer and "n" for an autoregressive flow on the 2-sphere. The autoregressive connectivity and amortization (for conditional PDFs) is taken care of by the module, and in the configuration indicated by a "+" connecting the different options.

Without much tuning, you should get something that just works and has the properties of a flexible PDF... however high customization of flow parameters and connectivity is also possible if desired.

It implements a few state of the art flows that to my knowledge are not really found in other repositories yet (e.g. Gaussianization flows).

Currently implemented manifolds with respective flows (from the README):

Euclidean flows:

• Generic affine flow (Multivariate normal distribution) ("t")
• Gaussianization flow arXiv:2003.01941 ("g")
• Hybrid of nonlinear scalings and rotations ("Polynomial Stretch flow") ("p")

1-sphere Flows:

• Moebius transformations (described in arXiv:2002.02428) ("m")
• Circular rational-quadratic splines (described in arXiv:2002.02428) ("o")

2-sphere Flows:

• Autorregressive flow for N-Spheres (arXiv:2002.02428) ("n")
• Exponential map flow (arXiv:0906.0874/arXiv:2002.02428) ("v")
• Neural Manifold Ordinary Differential Equations arXiv:2006.10254 ("c")

Interval Flows:

• "Neural Spline Flows" (Rational-quadratic splines) arXiv:1906.04032 ("r")

Simplex Flows:

• Autoregressive simplex flow arXiv:2008.05456 ("w")

All of those can be combined in any way and the package automatically manages the connectivity.
More info can also be found in the docs.

Best,
Thorsten

[Juan Emmanuel Johnson]

Hi, this looks great! I look forward to digging deeper into the package!

So you can make a pull request yourself and someone can merge it. You just need to add your package to the file data/packages.yml. Something like:

- id: pkg-11
  date: todays_date
  title: name_of_package
  url: package_website
  authors: ___
  authors_url: your_website
  lang: Python
  description: A NF package ....

Everything else should be handled by the .pre-commit-config.yaml and it should go smoothly! But either way, let us know if you have any other issues!

[Janosh Riebesell]

@thoglu Like @jejjohnson said, this would make a great addition! 👍 A PR to data/packages.yml would be great.