GenSBI: Generative Methods for Simulation-Based Inference in JAX

Overview

GenSBI is a JAX-based library for Simulation-Based Inference (SBI) using state-of-the-art generative methods, currently revolving around Optimal Transport Flow Matching, Diffusion Models, and Normalizing Flows.

It is designed for researchers and practitioners who need a flexible, high-performance toolkit to solve complex inference problems where the likelihood function is intractable.

The companion paper with a review of methodologies and functionalities is available at: https://arxiv.org/abs/2605.27499

Key Features

• Modern SBI Algorithms: Implements cutting-edge techniques like Optimal Transport Conditional Flow Matching, Diffusion Models, and Normalizing Flows for robust and flexible posterior inference.
• Built on JAX and Flax NNX: Leverages the power of JAX for automatic differentiation, vectorization, and seamless execution on CPUs, GPUs, and TPUs.
• High-Level Recipes API: A simplified interface for common workflows, allowing you to train models and run inference with just a few lines of code.
• Powerful Transformer Models: Includes implementations of recent, high-performing models like Flux1, Flux1Join, and Simformer for handling complex, high-dimensional data.
• Multi-Device Training and Sampling: Pass a mesh to any pipeline (from gensbi.sharding import make_mesh) for data-parallel training and sampling across GPUs/TPUs, with experimental FSDP and tensor-parallel rules.
• Modular and Extensible: A clean, well-structured codebase that is easy to understand, modify, and extend for your own research.

Installation

Using uv (recommended):

uv add gensbi
# or, for a standalone install:
uv pip install gensbi

For GPU support:

uv add gensbi[cuda12]
# or
uv pip install gensbi[cuda12]

Or using pip:

pip install gensbi

For GPU support and other options, including how to install uv, see the Installation Guide.

Quick Start

To get started immediately, you can use the high-level API to train a model.

from flax import nnx
from gensbi.recipes import Flux1FlowPipeline
from gensbi.models import Flux1Params

train_dataset = ... # define a training dataset (infinite iterator)
val_dataset = ...   # define a validation dataset (infinite iterator)
dim_obs = ...       # dimension of the parameters (theta)
dim_cond = ...      # dimension of the simulator observations (x)
params = Flux1Params(...) # the parameters for your model

# Instantiate the pipeline
pipeline = Flux1FlowPipeline(
    train_dataset,
    val_dataset,
    dim_obs,
    dim_cond,
    params=params,
)

# Train the model
# Note: GenSBI uses Flax NNX, so we pass a random key generator
pipeline.train(rngs=nnx.Rngs(0))

# After training, get a sampler for posterior sampling
key = jax.random.PRNGKey(42)
samples = pipeline.sample(key, x_observed, num_samples=10_000)

Examples

Examples for this library are available separately in the GenSBI-examples repository.

Some key examples include:

Getting Started:

• my_first_model.ipynb
  A beginner-friendly notebook introducing the core concepts of GenSBI on a simple problem.

Unconditional Density Estimation:

• flow_matching_2d_unconditional.ipynb
  Demonstrates how to use flow matching in 2D for unconditional density estimation.
• diffusion_EDM_2d_unconditional.ipynb
  Demonstrates how to use EDM diffusion models in 2D for unconditional density estimation.
• diffusion_SM_2d_unconditional.ipynb
  Demonstrates how to use score-matching diffusion models in 2D for unconditional density estimation.

Conditional Density Estimation:

• two_moons_flow_flux.ipynb
  Uses the Flux1 model for posterior density estimation on the two-moons benchmark.
• two_moons_diffusion_flux.ipynb
  Uses the Diffusion model for posterior density estimation on the two-moons benchmark.

Citing GenSBI

If you use this library, please consider citing this work and the original methodology papers, see references.

For the paper:

@article{amerio2026gensbi_paper,
      title={GenSBI: Generative Methods for Simulation-Based Inference in JAX}, 
      author={Aurelio Amerio},
      year={2026},
      eprint={2605.27499},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2605.27499}, 
}

For the software:

@software{amerio2026gensbi_software,
  author       = {Amerio, Aurelio},
  title        = {GenSBI: Generative Methods for Simulation-Based Inference in JAX
                  },
  month        = may,
  year         = 2026,
  publisher    = {Zenodo},
  version      = {v0.3.4},
  doi          = {10.5281/zenodo.20410084},
  url          = {https://doi.org/10.5281/zenodo.20410084},
}

Reference implementations

• Facebook Flow Matching library: https://github.com/facebookresearch/flow_matching
• Elucidating the Design Space of Diffusion-Based Generative Models: https://github.com/NVlabs/edm
• Simformer model: https://github.com/mackelab/simformer
• Flux1 model from BlackForest Lab: https://github.com/black-forest-labs/flux
• Simulation-Based Inference Benchmark: https://github.com/sbi-benchmark/sbibm