GitHub - femtomc/genjax: Probabilistic programming with vectorized programmable inference

Purpose: Vectorized probabilistic programming with generative functions and programmable inference in JAX.
Formal foundations: A mechanized Lean model proves product density preservation and fundamental vectorization for finite models under explicit primitive conformance certificates. Concrete execution maps vectorized generative primitives to XLA kernels, separating denotational density semantics from machine array compilation.
POPL 2026 artifact: v1.0.10

Use

git clone https://github.com/a-tiny-project/genjax.git
cd genjax
pixi install
pixi run test-fast
pixi run paper-figures
pixi run paper-figures-gpu

from genjax import gen, normal

@gen
def model():
    return normal(0.0, 1.0) @ "x"

trace = model.simulate()
choices = trace.get_choices()

Generative functions expose simulate, generate, assess, and update.
vmap and modular_vmap lift model and inference structure over explicit array axes.
Inspect all Pixi tasks in pyproject.toml.

Case	Figures	Command
Fair coin	16a	`pixi run paper-faircoin-gen`
Curve fitting	4–6	`pixi run paper-curvefit-gen`
Multi-framework benchmark	16b	`pixi run paper-perfbench`
Game of Life	18	`pixi run assets && pixi run -e gol gol-paper`
Localization	19	`pixi run paper-localization-gen`
AIR estimators	PLDI 2024 port	`pixi run air-compare`

Add --mode cuda to paper-perfbench for its CUDA pipeline.
CPU and GPU execute the same models but have different scaling curves.
Figure 19 and paper-scale curve fitting require CUDA-like throughput to match the published timing/ESS panels.
Gen.jl benchmark lanes require Julia 1.10 or newer.
Generated figures are saved in figs/. Perfbench outputs to separate CPU and CUDA directories.

Apache-2.0. See LICENSE.