Welcome to JAXNS’s documentation!
JAXNS is a probabilistic programming framework and advanced nested sampling algorithm. It’s goal is to empower researchers and scientists of all types, from early career to seasoned professionals, from small jupyter notebooks to massive HPC problem. Initially, I developed JAXNS to solve my own problems during my PhD. However, it has since grown into a full-fledged probabilistic programming framework. JAXNS has been applied in numerous domains from cosmology, astrophysics, gravitational waves, interferometry, exoplanets, particle physics, meta materials, epidemiology, climate modelling, and beyond. Not to mention it has been used in industry for a variety of applications. All of this is welcomed and gladly supported.
For a deeper understanding of the JAXNS v3 algorithm, read and cite Phantom-Conditioned Nested Sampling (Albert, 2026). It describes the race-tree formulation, phantom conditioning, and evidence-improving and posterior-improving allocation.
The previous Phantom-Powered Nested Sampling paper is redacted. The original JAXNS paper remains valid as a reference for JAXNS’s high performance, but its algorithm description is superseded by the v3 paper. See Papers and citation for the references and citation details.
Here are nine things you can do with JAXNS:
Build probabilistic models in an easy to use, high-level language, that can be used anywhere in the JAX ecosystem.
Compute the Bayesian evidence of a model or hypothesis (the ultimate scientific method);
Produce high-quality samples from the posterior distribution;
Easily handle degenerate difficult multi-modal posteriors;
Model both discrete and continuous priors;
Encode complex constraints on the prior space;
Easily embed your neural networks or ML model in the likelihood/prior;
Easily embed JAXNS in your ML model;
Use JAXNS in a distributed computing environment;
JAXNS’s Mission Statement
Our mission is to make nested sampling faster, easier, and more powerful.
User Guide
API Reference
Examples
- Inference of Jones scalar observables (noisy angular quantities)
- Dark-matter detector mixture
- Dual-moon likelihood
- Poisson likelihood and Gamma prior
- Gaussian-process kernel comparison on contaminated data
- Thin Gaussian Shells with Uniform Prior
- Linear regression with additional scatter
- Multivariate Normal Likelihood with Multivariate Normal Prior
- OU process
- Inference from a noisy logical report
- Self-Exciting process (Hawkes process)
- Polynomial-order inference with a fixed-dimensional product space