Neural simulation-based inference (SBI) describes an emerging family of methods for Bayesian inference with intractable likelihood functions that use neural networks as surrogate models. Here we introduce sbijax, a Python package that implements a wide variety of state-of-the-art methods in neural simulation-based inference using a user-friendly programming interface. sbijax offers high-level functionality to quickly construct SBI estimators, and compute and visualize posterior distributions with only a few lines of code. In addition, the package provides functionality for conventional approximate Bayesian computation, to compute model diagnostics, and to automatically estimate summary statistics. By virtue of being entirely written in JAX, sbijax is extremely computationally efficient, allowing rapid training of neural networks and executing code automatically in parallel on both CPU and GPU.

翻译：神经仿真推断（SBI）描述了一类新兴的贝叶斯推断方法，这些方法针对似然函数难以处理的问题，利用神经网络作为代理模型。本文介绍sbijax——一个Python软件包，它通过用户友好的编程接口实现了多种最先进的神经仿真推断方法。sbijax提供高级功能，仅需少量代码即可快速构建SBI估计器、计算并可视化后验分布。此外，该软件包还支持传统的近似贝叶斯计算、模型诊断计算以及自动估计摘要统计量。由于完全基于JAX编写，sbijax具有极高的计算效率，能够快速训练神经网络，并自动在CPU和GPU上并行执行代码。