mlx-examples/flow
2023-12-18 00:48:25 -05:00
..
bijectors.py Implement normalizing flow Real NVP example 2023-12-18 00:48:25 -05:00
distributions.py Implement normalizing flow Real NVP example 2023-12-18 00:48:25 -05:00
flows.py Implement normalizing flow Real NVP example 2023-12-18 00:48:25 -05:00
main.py Implement normalizing flow Real NVP example 2023-12-18 00:48:25 -05:00
README.md Implement normalizing flow Real NVP example 2023-12-18 00:48:25 -05:00
samples.png Implement normalizing flow Real NVP example 2023-12-18 00:48:25 -05:00

Normalizing flow

Real NVP normalizing flow from Dinh et al. (2016) implemented using mlx.

The example is written in a somewhat more object-oriented style than strictly necessary, with an eye towards extension to other use cases benefitting from arbitrary distributions and bijectors.

Usage

The example can be run with

python main.py

which trains the normalizing flow on the two moons dataset and plots the result in samples.png.

By default the example runs on the GPU. To run on the CPU, do

python main.py --cpu

For all available options, run

python main.py --help

Results

Samples