Probabilistic Programming over Heterogeneous Language and Hardware Targets

Abstract

Modern probabilistic programming applications, from large-scale Bayesian inference to real-time decision making, require both the expressiveness of CPU-oriented languages such as Gen.jl and the massive parallelism of GPU-backed array languages such as GenJAX, yet existing platforms force users to trade modeling flexibility for performance. This thesis introduces GenUflect, a metalanguage that embeds multiple Gen-compatible dialects inside a single program, allowing each sub-component to run on the most appropriate language and hardware target while preserving Gen’s programmable-inference interface. GenUflect extends Gen’s dynamic-modeling language with the @union, @vmap, @amortize, @amortize≤, and @runtime_union combinators; these macros compile at build-time (or justin-time) to autonomous generative functions written in the target dialect, link them through a lightweight FFI layer, and manage cross-device data via zero-copy MirrorArrays and lazily materialized traces. The resulting programs remain sound by construction because each foreign subtrace is itself a valid Gen generative function. Empirical studies demonstrate that this hybrid approach yields large practical gains. On a split linear-vs-sinusoidal regression task, GenUflect matches pure GenJAX throughput while running higher-order control logic on the CPU, and is up to two orders of magnitude faster than a pure Gen implementation for datasets of 105 points. In a collapsed-Gibbs sampler for a Dirichlet-process mixture model, GenUflect’s elastic allocation (@amortize≤) lets vectorized GPU kernels adapt to a growing number of clusters; the same inference that takes over an hour in Gen executes in seconds with GenUflect. A probabilistic inverse-graphics pipeline further showcases how heterogeneous submodels can cooperate seamlessly within unified inference code. By coupling language interoperability with automated data movement and compile-time code generation, GenUflect bridges the gap between flexibility and speed, enabling scalable, expressive probabilistic programs that natively exploit both CPUs and accelerators.