Probabilistic Programming with Vectorized Programmable Inference
summary
GenJAX is a Python-embedded language and compiler that integrates JAX and exposes a vmap vectorization transformation to vectorize probabilistic traces, stochastic branching and inference interfaces. The paper formalizes vectorization on a core lambda_GEN calculus and proves correctness, and the implementation compiles to JAX/ XLA. Benchmarks show performance competitive with hand-optimized JAX; case studies include 512×512 inversion in milliseconds for Conway’s Game of Life and a robot-localization SMC variant that tracks reliably in ~2 ms per 16-step run.
AI-written summary · about
chapters (13)
- 0:00Introduction
- 1:44Vectorized inference design
- 3:29VMAP for parallelism
- 7:01Generative functions syntax
- 8:29Probabilistic interfaces
- 10:54Compiler transformations
- 12:22Vectorization syntax
- 14:06Nesting VMAP
- 16:12Custom inference with VMAP
- 17:58GenJAX compiler
- 18:59Conway's Game of Life inversion
- 20:23Robot localization SMC
- 22:06Conclusion
related talks
11mIntro to JAX: Accelerating Machine Learning research
01h 10mJAX: accelerated machine learning research via composable function transformations in Python
31mHigh performance machine learning with JAX
14mJAX: Accelerated Machine Learning Research via Composable Function Transformations in Python
41mJAX for Bayes
02h 24mMagical NumPy with JAX
21mProbabilistic Programming with Programmable Variational Inference
16mA Domain-Specific Probabilistic Programming Language for Reasoning about Reasoning(…)