├── .github └── workflows │ ├── docs.yml │ ├── publish.yml │ └── test.yml ├── .gitignore ├── LICENSE ├── README.md ├── docs ├── Makefile ├── README.md ├── _static │ ├── icon.svg │ ├── logo_dark.svg │ ├── logo_light.svg │ └── style.css ├── api │ ├── bijections.rst │ ├── distributions.rst │ ├── experimental.rst │ ├── flows.rst │ ├── losses.rst │ └── training.rst ├── conf.py ├── examples │ ├── conditional.ipynb │ ├── constrained.ipynb │ ├── examples.rst │ ├── snpe.ipynb │ ├── unconditional.ipynb │ └── variational_inference.ipynb ├── faq.rst ├── getting_started.rst ├── index.rst └── make.bat ├── flowjax ├── __init__.py ├── bijections │ ├── __init__.py │ ├── affine.py │ ├── bijection.py │ ├── block_autoregressive_network.py │ ├── chain.py │ ├── concatenate.py │ ├── coupling.py │ ├── exp.py │ ├── jax_transforms.py │ ├── masked_autoregressive.py │ ├── orthogonal.py │ ├── planar.py │ ├── power.py │ ├── rational_quadratic_spline.py │ ├── sigmoid.py │ ├── softplus.py │ ├── tanh.py │ └── utils.py ├── distributions.py ├── experimental │ ├── __init__.py │ └── numpyro.py ├── flows.py ├── masks.py ├── py.typed ├── root_finding.py ├── tasks.py ├── train │ ├── __init__.py │ ├── loops.py │ ├── losses.py │ └── train_utils.py └── utils.py ├── pyproject.toml └── tests ├── test_bijections ├── test_bijection_utils.py ├── test_bijections.py ├── test_bnaf.py ├── test_chain.py ├── test_jax_transforms.py ├── test_masked_autoregressive.py └── test_rational_quadratic_spline.py ├── test_distributions.py ├── test_experimental └── test_numpyro.py ├── test_flows.py ├── test_masks.py ├── test_root_finding.py ├── test_train ├── test_loops.py ├── test_losses.py └── test_train_utils.py └── test_utils.py /.github/workflows/docs.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/.github/workflows/docs.yml -------------------------------------------------------------------------------- /.github/workflows/publish.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/.github/workflows/publish.yml -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/.github/workflows/test.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/README.md -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/Makefile -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/README.md -------------------------------------------------------------------------------- /docs/_static/icon.svg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/_static/icon.svg -------------------------------------------------------------------------------- /docs/_static/logo_dark.svg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/_static/logo_dark.svg -------------------------------------------------------------------------------- /docs/_static/logo_light.svg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/_static/logo_light.svg -------------------------------------------------------------------------------- /docs/_static/style.css: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/_static/style.css -------------------------------------------------------------------------------- /docs/api/bijections.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/api/bijections.rst -------------------------------------------------------------------------------- /docs/api/distributions.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/api/distributions.rst -------------------------------------------------------------------------------- /docs/api/experimental.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/api/experimental.rst -------------------------------------------------------------------------------- /docs/api/flows.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/api/flows.rst -------------------------------------------------------------------------------- /docs/api/losses.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/api/losses.rst -------------------------------------------------------------------------------- /docs/api/training.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/api/training.rst -------------------------------------------------------------------------------- /docs/conf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/conf.py -------------------------------------------------------------------------------- /docs/examples/conditional.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/examples/conditional.ipynb -------------------------------------------------------------------------------- /docs/examples/constrained.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/examples/constrained.ipynb -------------------------------------------------------------------------------- /docs/examples/examples.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/examples/examples.rst -------------------------------------------------------------------------------- /docs/examples/snpe.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/examples/snpe.ipynb -------------------------------------------------------------------------------- /docs/examples/unconditional.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/examples/unconditional.ipynb -------------------------------------------------------------------------------- /docs/examples/variational_inference.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/examples/variational_inference.ipynb -------------------------------------------------------------------------------- /docs/faq.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/faq.rst -------------------------------------------------------------------------------- /docs/getting_started.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/getting_started.rst -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/index.rst -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/docs/make.bat -------------------------------------------------------------------------------- /flowjax/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/__init__.py -------------------------------------------------------------------------------- /flowjax/bijections/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/__init__.py -------------------------------------------------------------------------------- /flowjax/bijections/affine.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/affine.py -------------------------------------------------------------------------------- /flowjax/bijections/bijection.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/bijection.py -------------------------------------------------------------------------------- /flowjax/bijections/block_autoregressive_network.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/block_autoregressive_network.py -------------------------------------------------------------------------------- /flowjax/bijections/chain.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/chain.py -------------------------------------------------------------------------------- /flowjax/bijections/concatenate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/concatenate.py -------------------------------------------------------------------------------- /flowjax/bijections/coupling.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/coupling.py -------------------------------------------------------------------------------- /flowjax/bijections/exp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/exp.py -------------------------------------------------------------------------------- /flowjax/bijections/jax_transforms.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/jax_transforms.py -------------------------------------------------------------------------------- /flowjax/bijections/masked_autoregressive.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/masked_autoregressive.py -------------------------------------------------------------------------------- /flowjax/bijections/orthogonal.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/orthogonal.py -------------------------------------------------------------------------------- /flowjax/bijections/planar.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/planar.py -------------------------------------------------------------------------------- /flowjax/bijections/power.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/power.py -------------------------------------------------------------------------------- /flowjax/bijections/rational_quadratic_spline.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/rational_quadratic_spline.py -------------------------------------------------------------------------------- /flowjax/bijections/sigmoid.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/sigmoid.py -------------------------------------------------------------------------------- /flowjax/bijections/softplus.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/softplus.py -------------------------------------------------------------------------------- /flowjax/bijections/tanh.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/tanh.py -------------------------------------------------------------------------------- /flowjax/bijections/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/bijections/utils.py -------------------------------------------------------------------------------- /flowjax/distributions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/distributions.py -------------------------------------------------------------------------------- /flowjax/experimental/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /flowjax/experimental/numpyro.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/experimental/numpyro.py -------------------------------------------------------------------------------- /flowjax/flows.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/flows.py -------------------------------------------------------------------------------- /flowjax/masks.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/masks.py -------------------------------------------------------------------------------- /flowjax/py.typed: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /flowjax/root_finding.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/root_finding.py -------------------------------------------------------------------------------- /flowjax/tasks.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/tasks.py -------------------------------------------------------------------------------- /flowjax/train/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/train/__init__.py -------------------------------------------------------------------------------- /flowjax/train/loops.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/train/loops.py -------------------------------------------------------------------------------- /flowjax/train/losses.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/train/losses.py -------------------------------------------------------------------------------- /flowjax/train/train_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/train/train_utils.py -------------------------------------------------------------------------------- /flowjax/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/flowjax/utils.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/pyproject.toml -------------------------------------------------------------------------------- /tests/test_bijections/test_bijection_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_bijections/test_bijection_utils.py -------------------------------------------------------------------------------- /tests/test_bijections/test_bijections.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_bijections/test_bijections.py -------------------------------------------------------------------------------- /tests/test_bijections/test_bnaf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_bijections/test_bnaf.py -------------------------------------------------------------------------------- /tests/test_bijections/test_chain.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_bijections/test_chain.py -------------------------------------------------------------------------------- /tests/test_bijections/test_jax_transforms.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_bijections/test_jax_transforms.py -------------------------------------------------------------------------------- /tests/test_bijections/test_masked_autoregressive.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_bijections/test_masked_autoregressive.py -------------------------------------------------------------------------------- /tests/test_bijections/test_rational_quadratic_spline.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_bijections/test_rational_quadratic_spline.py -------------------------------------------------------------------------------- /tests/test_distributions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_distributions.py -------------------------------------------------------------------------------- /tests/test_experimental/test_numpyro.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_experimental/test_numpyro.py -------------------------------------------------------------------------------- /tests/test_flows.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_flows.py -------------------------------------------------------------------------------- /tests/test_masks.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_masks.py -------------------------------------------------------------------------------- /tests/test_root_finding.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_root_finding.py -------------------------------------------------------------------------------- /tests/test_train/test_loops.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_train/test_loops.py -------------------------------------------------------------------------------- /tests/test_train/test_losses.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_train/test_losses.py -------------------------------------------------------------------------------- /tests/test_train/test_train_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_train/test_train_utils.py -------------------------------------------------------------------------------- /tests/test_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/danielward27/flowjax/HEAD/tests/test_utils.py --------------------------------------------------------------------------------