Pus$\mathbb{H}$: Concurrent Probabilistic Programming with Function Spaces

We introduce a prototype probabilistic programming language (PPL) called Pus$\mathbb{H}$ for performing Bayesian inference on function spaces with a focus on Bayesian deep learning (BDL). We describe the core abstraction of Pus$\mathbb{H}$ based on particles that links models, specified as neural networks (NNs), with inference, specified as procedures on particles using a programming model inspired by message passing. Finally, we test Pus$\mathbb{H}$ on a variety of models and datasets used in scientific machine learning (SciML), a domain with natural function space inference problems, and we evaluate scaling of Pus$\mathbb{H}$ on single-node multi-GPU devices. Thus we explore the combination of probabilistic programming, NNs, and concurrency in the context of Bayesian inference on function spaces. The code can be found at https://github.com/lbai-lab/PusH.