Probabilistic Programming with Vectorized Programmable Inference
Author(s)
Becker, McCoy R.
Huot, Mathieu
Matheos, George
Wang, Xiaoyan
Chung, Karen
Smith, Colin
Ritchie, Sam
Saurous, Rif A.
Lew, Alexander K.
Rinard, Martin C.
Date Issued
January 8, 2026
Journal
Proceedings of the ACM on Programming Languages
Publisher
ACM
Citation
McCoy R. Becker, Mathieu Huot, George Matheos, Xiaoyan Wang, Karen Chung, Colin Smith, Sam Ritchie, Rif A. Saurous, Alexander K. Lew, Martin C. Rinard, and Vikash K. Mansinghka. 2026. Probabilistic Programming with Vectorized Programmable Inference. Proc. ACM Program. Lang. 10, POPL, Article 87 (January 2026), 32 pages.
Version
Final published version
Abstract
We present GenJAX, a new language and compiler for vectorized programmable probabilistic inference.
GenJAX integrates the vectorizing map (vmap) operation from array programming frameworks such as JAX
into the programmable inference paradigm, enabling compositional
vectorization of features such as probabilistic program traces, stochastic branching
(for expressing mixture models), and programmable inference interfaces
for writing custom probabilistic inference algorithms.
We formalize vectorization as a source-to-source program transformation on a core calculus for probabilistic programming ($\gen$), and
prove that it correctly vectorizes both modeling and inference operations.
We have implemented our approach in \href{https://github.com/probcomp/genjax}{the GenJAX language and compiler}, and have empirically evaluated this implementation on
several benchmarks and case studies. Our results show that our implementation
supports a wide and expressive set of programmable inference patterns and delivers
performance comparable to hand-optimized JAX code.
MIT Department
Massachusetts Institute of Technology. Department of Electrical Engineering and Computer Science
Massachusetts Institute of Technology. Department of Brain and Cognitive Sciences
Terms of Use
Creative Commons Attribution
Persistent DSpace Link
DOI of Published Version
https://doi.org/10.1145/3776729
