ProbJax¶

A JAX-native research toolbox for probabilistic computation: probabilistic-program transformations, SciPy-style distributions, automatic program inversion, a large inference suite, and Flax NNX generative models. Core array computations compose with JAX transformations; host callbacks, data loaders and some solver gradient paths have additional restrictions described in their guides.

Note

ProbJax is research software. APIs evolve, and the accelerator-specific kernels and sharding features are experimental.

import jax
import jax.numpy as jnp
from probjax.core import condition, joint_sample, log_joint_fn
from probjax.stats import norm

def model(key):
    key_z, key_y = jax.random.split(key)
    z = norm.rvs(key_z, 0.0, 1.0, name="z")
    return norm.rvs(key_y, z, 0.5, name="y")

# Sample every named site at once
sites = joint_sample(model)(jax.random.key(0))

# Condition on an observation and score the latent
posterior = condition(model, {"y": jnp.asarray(0.25)})
latent = joint_sample(posterior)(jax.random.key(1))["z"]
log_joint = log_joint_fn(posterior)(z=latent)

What is here¶

Probabilistic programs

Write an ordinary JAX function, then trace, condition, observe, intervene and joint_sample it. Sampling a distribution inside a transformation records a named site; outside one it is a plain JAX sample.
Program inversion

inverse and inverse_and_logabsdet walk a jaxpr backwards to build the inverse of a function, with custom_inverse for the parts that need an analytic rule. See Program inversion.
Density estimation

Normalizing flows, autoregressive models, diffusion and flow matching, with one fit and one as_dist interface across all of them. See Density estimation.
Inference

MCMC kernels, SMC, Kalman-family filters and particle filtering, plus flow-based variational inference and NeuTra preconditioning. See Inference.

New here?¶

Start with the Overview for how the modules fit together and the Glossary for the vocabulary. Then follow Getting started; if something breaks, check Troubleshooting, and to work on ProbJax itself see Contributing.

Install¶

git clone https://github.com/mackelab/probjax.git
cd probjax
python -m pip install -e .

These pages track main; the published package can lag behind them. For the next release, see migration notes.

Python 3.11–3.13 and JAX >=0.9.0,<0.12. Accelerator options are in Getting started.

Tutorials¶

A curated set of example notebooks is rendered as part of these pages (generated from examples/ without re-execution, so outputs are those committed in the notebooks; CI executes the generated Python examples, but does not regenerate their displayed outputs):

All runnable notebooks live in examples/ (core, stats, inference, nn, utils).

The remaining notebooks may predate recent API refactors. The curated tutorial pages have executable examples checked by CI; stored plots and other outputs may reflect an earlier execution. Prefer the guides and reference when they disagree with an older notebook.