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上并行执行代码。