{"id":40091145,"url":"https://github.com/mathisgerdes/bijx","last_synced_at":"2026-01-19T10:03:46.811Z","repository":{"id":311601101,"uuid":"1029789636","full_name":"mathisgerdes/bijx","owner":"mathisgerdes","description":"Bijections \u0026 normalizing flows with JAX/NNX","archived":false,"fork":false,"pushed_at":"2025-12-08T09:14:32.000Z","size":2646,"stargazers_count":11,"open_issues_count":0,"forks_count":0,"subscribers_count":1,"default_branch":"master","last_synced_at":"2025-12-10T06:26:21.478Z","etag":null,"topics":["bijections","flax","jax","lattice-field-theory","lie-groups","nnx","normalizing-flows","sampling"],"latest_commit_sha":null,"homepage":"https://www.mathisgerdes.com/bijx/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/mathisgerdes.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":"CITATION.cff","codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2025-07-31T15:18:22.000Z","updated_at":"2025-12-09T05:48:56.000Z","dependencies_parsed_at":"2025-08-25T13:27:43.180Z","dependency_job_id":"41988774-55e3-4c7a-8e8a-264ee9b86180","html_url":"https://github.com/mathisgerdes/bijx","commit_stats":null,"previous_names":["mathisgerdes/bijx"],"tags_count":6,"template":false,"template_full_name":null,"purl":"pkg:github/mathisgerdes/bijx","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mathisgerdes%2Fbijx","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mathisgerdes%2Fbijx/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mathisgerdes%2Fbijx/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mathisgerdes%2Fbijx/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/mathisgerdes","download_url":"https://codeload.github.com/mathisgerdes/bijx/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mathisgerdes%2Fbijx/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":28565056,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-01-19T08:53:44.001Z","status":"ssl_error","status_checked_at":"2026-01-19T08:52:40.245Z","response_time":67,"last_error":"SSL_connect returned=1 errno=0 peeraddr=140.82.121.6:443 state=error: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["bijections","flax","jax","lattice-field-theory","lie-groups","nnx","normalizing-flows","sampling"],"created_at":"2026-01-19T10:03:46.749Z","updated_at":"2026-01-19T10:03:46.804Z","avatar_url":"https://github.com/mathisgerdes.png","language":"Python","readme":"\u003cp align=\"center\"\u003e\n    \u003cimg src=\"https://raw.githubusercontent.com/mathisgerdes/bijx/master/docs/source/_static/icons/bijx.svg\" alt=\"bijx logo\" height=\"100\"\u003e\n    \u003cbr\u003e\n    \u003cem\u003e\u003cb\u003ebij\u003c/b\u003eection + ja\u003cb\u003ex\u003c/b\u003e = /ˈbaɪdʒæks/\u003c/em\u003e\n\u003c/p\u003e\n\n# Bijections \u0026 normalizing flows with JAX/NNX\n[![GitHub](https://img.shields.io/badge/code-gray?logo=github\u0026logoColor=white)](https://github.com/mathisgerdes/bijx)\n[![docs build](https://github.com/mathisgerdes/bijx/actions/workflows/docs.yml/badge.svg)](https://mathisgerdes.github.io/bijx/)\n[![Tests](https://github.com/mathisgerdes/bijx/actions/workflows/tests.yml/badge.svg)](https://github.com/mathisgerdes/bijx/actions/workflows/tests.yml)\n[![codecov](https://codecov.io/github/mathisgerdes/bijx/graph/badge.svg?token=QYAJU8QCJE)](https://codecov.io/github/mathisgerdes/bijx)\n[![PyPI](https://img.shields.io/pypi/v/bijx)](https://pypi.org/project/bijx/)\n[![DOI](https://joss.theoj.org/papers/10.21105/joss.09521/status.svg)](https://doi.org/10.21105/joss.09521)\n\nThis library provides flexible tools for building normalizing flows and bijections with tractable change of densities, focusing on research and applications in physics. It aims to provide reusable building blocks rather than a simplified interface for common use cases.\n\nThe library is built around two fundamental mathematical objects:\n\n- **Bijections**: Invertible transformations that track their effect on probability densities.\n- **Distributions**: Probability distributions with methods for sampling and density evaluation.\n\n## Key Features\n\nCore design goals:\n- **Modular Building Blocks**: Composable components rather than monolithic frameworks\n- **Research-Focused**: Prioritizes flexibility and expressiveness over convenience\n- **Flax NNX Integration**: Modern state management for neural network components\n\nPhysics \u0026 advanced methods:\n- **Continuous Normalizing Flows**: CNFs as core primitive with flexible ODE solver backends\n- **Matrix Lie Group Operations**: Automatic differentiation on matrix groups\n- **Structure-Preserving Integration**: Crouch-Grossmann solvers for ODEs on matrix groups\n- **Symmetry-Aware Architectures**: Equivariant layers and transformations (e.g., lattice-symmetric CNFs)\n- **Fourier-Space Operations**: Tools for momentum space transformations and complex field decomposition\n\n## Quickstart\n\nHere is a minimal example of building and sampling from a simple normalizing flow. We transform samples from a base distribution (a standard normal) using a chain of bijections to produce samples from a new, transformed distribution.\n\n```python\nimport jax\nimport jax.numpy as jnp\nimport bijx\nfrom flax import nnx\n\n# Define simple base distributions\nprior = bijx.IndependentNormal(event_shape=(2,))\n\n# choose a base bijection for the coupling layer\nbase_bijection = bijx.ModuleReconstructor(\n    # for real NVP style, using affine linear as base bijection\n    bijx.AffineLinear(rngs=nnx.Rngs(0))\n)\n\n# build a simple coupling layer\ncoupling_layer = bijx.GeneralCouplingLayer(\n    # replace with more realistic NN (active -\u003e params)\n    embedding_net=lambda x: jnp.zeros(x.shape + (base_bijection.params_total_size,)),\n    mask=bijx.BinaryMask.from_boolean_mask(jnp.array([True, False])),\n    bijection_reconstructor=base_bijection,\n)\n\n# build a very simple continuous flow\n# note: automatic divergence is expensive in high dimensions\ncnf = bijx.ContFlowDiffrax(\n    # vf must return divergence too; AutoJacVF computes it for us\n    vf=bijx.AutoJacVF(\n        lambda t, x: -x,\n    )\n)\n\n# Compose bijections\nflow = bijx.Chain(\n    bijx.Shift(jnp.array([5.0, -2.0])),\n    bijx.Scaling(jnp.array([2.0, 0.5])),\n    # add more coupling layers for expressivity\n    coupling_layer,\n    # final continuous flow\n    cnf,\n)\n\n# sample from the base distribution\nrng = jax.random.PRNGKey(42)\nx, lp_x = prior.sample(batch_shape=(5,), rng=rng)\n\n# transform samples with bijection\ny, lp_y = flow.forward(x, lp_x)\n\n# these can also be generated by combining prior and flow into a new distribution\ndist = bijx.Transformed(prior, flow)\n\ny2, lp_y2 = dist.sample(batch_shape=(5,), rng=rng)\n\n# exactly the same as we used the same random key\nassert jnp.allclose(y, y2)\n\n# check bijectivity\nx_rec, lp_x_rec = flow.reverse(y, lp_y)\nassert jnp.allclose(x, x_rec, atol=1e-6)\nassert jnp.allclose(lp_x, lp_x_rec, atol=1e-6)\n```\n\n## Design Principles\n\n- **Modularity**: Different parts of the library should be usable on their own. Expose as many expressive building blocks as possible.\n- **Flexibility**: Prioritize flexibility, sometimes at the cost of adding more ways to break things.\n- **Runtime Shape Inference**: Use a `batch + space + channels` convention and automatic vectorization for flexible data shape handling.\n\n## Installation\n\nThis package can be installed directly from PyPI:\n\n```bash\npip install bijx\n```\n\nAlternatively, to install locally from source:\n\n```bash\npip install -e .\n```\n\nFor development, install as an editable package with all dependencies (optionally without `docs` if documentation building is not needed):\n\n```bash\npip install -e \".[dev,docs]\"\n```\n\nTo keep the codebase tidy, please install `pip install pre-commit` and run `pre-commit install` before committing changes.\n\n## Documentation\n\nTo compile and open a local server for the documentation, run `make livehtml` in the `docs/` directory.\n\n## Testing\n\nThere are two types of tests that can be run:\nthe unit tests in `tests/` and the docstring examples in the source code `src/bijx/`.\n\n```bash\n# Run tests\npytest tests/\n\n# Run doctests\npytest src/bijx/ --doctest-modules\n\n# Run all tests\npytest tests/ src/bijx/ --doctest-modules\n\n# Run tests in parallel using xdist\npytest -n auto\n```\n\n## Citation\n\nIf you use bijx in your research, please cite our JOSS paper:\n\n```bibtex\n@article{Gerdes2025Bijx,\n  doi = {10.21105/joss.09521},\n  url = {https://doi.org/10.21105/joss.09521},\n  year = {2025},\n  publisher = {The Open Journal},\n  volume = {10},\n  number = {116},\n  pages = {9521},\n  author = {Gerdes, Mathis and Cheng, Miranda C. N.},\n  title = {Bijx: Bijections and normalizing flows with JAX/NNX},\n  journal = {Journal of Open Source Software}\n}\n```\n\n### Related Publications\n\nThis library extends research software developed for the following publications, which you may find useful and may consider citing if you use the relevant components:\n\n- **[Learning Lattice Quantum Field Theories with Equivariant Continuous Flows](https://arxiv.org/abs/2207.00283)**\n\n- **[Non-Perturbative Trivializing Flows for Lattice Gauge Theories](https://arxiv.org/abs/2410.13161)**\n\n- **[Analytic Bijections for Smooth and Interpretable Normalizing Flows](https://arxiv.org/abs/2601.10774)**\n\n## Module Layout\n\nThe library is organized into core mathematical tools, machine learning components, and specific applications.\n\n```\nbijx/\n├── __init__.py           # Main package exports\n├── utils.py              # General utilities\n├── distributions.py      # Core Distribution classes\n├── samplers.py           # Sampling helpers\n├── solvers.py            # ODE solver implementations\n│\n├── bijections/           # All bijection-related code\n│   ├── base.py           # Core Bijection, Chain, etc.\n│   ├── continuous.py     # Continuous flow wrappers\n│   ├── conv_cnf.py       # Convolutional CNF architecture\n│   ├── coupling.py       # Coupling layers\n│   └── ...\n│\n├── nn/                   # Neural network components (NNX)\n│   ├── conv.py           # Symmetric convolutions\n│   ├── embeddings.py     # Time and positional embeddings\n│   └── features.py       # Reusable feature-mappers\n│\n├── fourier.py            # Tools for Fourier-space operations\n├── lie.py                # General-purpose Lie group operations\n├── cg.py                 # Crouch-Grossmann ODE integrators\n│\n└── lattice/              # Application: Lattice Field Theory\n    ├── scalar.py         # Action and observables for phi^4 theory\n    └── gauge.py          # Gauge field symmetry operations\n```\n","funding_links":[],"categories":[],"sub_categories":[],"project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmathisgerdes%2Fbijx","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fmathisgerdes%2Fbijx","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmathisgerdes%2Fbijx/lists"}