Overview¶

ProbJax is a JAX-native library for probabilistic computation. It is organised into a small number of top-level modules, each doing one job: probabilistic programs and their transformations, SciPy-shaped distributions, Flax NNX generative models, an inference suite, and the numerical helpers the rest are built on.

This page is a map. Each section names the module, what lives in it, and where to read more. If a term is unfamiliar, the Glossary defines the vocabulary used throughout.

probjax.core — probabilistic programs¶

A model is an ordinary JAX function. When it draws a named value under a ProbJax transformation, the draw records a site; the transformations then rewire the program: condition on data (condition), intervene causally (do), fix values (substitute), or group auto-named draws (scope). Sampling and scoring use joint_sample, log_joint_fn and trace(..., sites=True) reports every site it saw. observe, intervene, and log_potential_fn are deprecated aliases.

The module also holds automatic program inversion — inverse, inverse_and_logabsdet and custom_inverse walk a jaxpr backwards, with JaxprGraph exposing the traced graph.

• Probabilistic programs — the transformation guide
• Program inversion — inverting a function
• Core reference — the full symbol list

probjax.stats — distributions¶

SciPy-shaped distributions: lowercase generators such as norm, gamma, dirichlet and categorical, used directly (norm.rvs(key, loc, scale)) or frozen with fixed parameters (norm(loc=..., scale=...)). Higher-order constructors — indep, mixture and transformed — build joint, mixture and pushed-forward distributions. The module also provides the training loop, fit, shared by the generative models.

• Distributions — using, freezing and fitting
• Stats reference — every family and method

probjax.nn — generative models¶

Flax NNX modules for density estimation: normalizing flows, autoregressive models, diffusion models and flow matching. They share one interface — construct with nnx.Rngs, call fit to train, and as_dist() for a frozen distribution with logpdf and sample.

• Density estimation — the shared interface
• Reference — per-family documentation

probjax.inference — inference¶

MCMC kernels, sequential Monte Carlo, Kalman-family filters, particle filtering and flow-based variational inference, with NeuTra preconditioning. The inference algorithms take a log-density callable (such as one built by log_joint_fn) and a key, with model-specific adapters where supported.

• Inference — MCMC and variational inference
• Sequential Monte Carlo and Temporal SMC — particle methods

probjax.utils — numerical building blocks¶

ODE and SDE integrators, root finders, special functions (inverse incomplete beta and gamma, log1mexp, logdiffexp), interpolation and linear-algebra helpers, and the shared type aliases used across the package.

• ODE and SDE integration — solvers and gradients
• Utilities reference — the full symbol list

Where to go next¶

• Getting started — install and first steps
• Glossary — the vocabulary these pages use
• Troubleshooting — common errors and their fixes
• Contributing — development workflow and documentation build