一句话：提出 GenJAX 语言与编译器，将 vmap 向量化引入可编程概率推断，并证明其模型与推断操作向量化的正确性。
本文面向概率编程中可编程推断难以高效向量化的问题，提出 GenJAX——一种支持向量化可编程概率推断的新语言与编译器。作者将数组编程框架（如 JAX）中的向量化映射 vmap 操作集成到可编程推断范式中，从而对概率程序轨迹、随机分支（用于表达混合模型）以及自定义推断算法接口等特性实现组合式向量化。方法上，作者在概率编程核心演算 λGEN 上将向量化形式化为源到源程序变换，并证明该变换能正确向量化建模与推断操作。作者实现了 GenJAX 语言与编译器，并在多个基准与案例研究上进行实证评估。结果表明，该实现支持广泛且富有表达力的可编程推断模式，并能达到与手写优化 JAX 代码相当的性能。
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 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.
@article{BeckerHMWCSRSLRM26,
title = {Probabilistic Programming with Vectorized Programmable Inference},
author = {McCoy R. Becker and Mathieu Huot and George Matheos and Xiaoyan Wang and Karen Chung and Colin Smith and Sam Ritchie and Rif A. Saurous and Alexander K. Lew and Martin C. Rinard and Vikash K. Mansinghka},
journal = {Proceedings of the ACM on Programming Languages},
volume = {10},
number = {POPL},
year = {2026},
doi = {10.1145/3776729}
}