gensbi

GenSBI: A library for Simulation-Based Inference (SBI) using Optimal Transport Flow Matching and Diffusion models in JAX.

Provides tools for probabilistic modeling, simulation, and training of generative models, including:

• Flow Matching techniques

• Diffusion models (EDM, score matching)

• Transformer-based models (Flux1, Simformer)

See the documentation for details and usage examples.

Submodules

• gensbi._compat
• gensbi.core
• gensbi.diagnostics
• gensbi.diffusion
• gensbi.experimental
• gensbi.flow_matching
• gensbi.inference
• gensbi.models
• gensbi.normalizing_flows
• gensbi.recipes
• gensbi.sharding
• gensbi.solver
• gensbi.utils

Attributes

__version__

Functions

_guard_grain_profiler_flag()

Let grain read its worker-profiling absl flag in a non-absl process.

Package Contents

gensbi._guard_grain_profiler_flag()

Let grain read its worker-profiling absl flag in a non-absl process.

grain 0.2.18 (grain/_src/core/profiler.py) reads --grain_enable_multiprocess_worker_profiling inside mp_prefetch whenever jaxlib_extension_version >= 448 (jax >= 0.11). absl raises UnparsedFlagAccessError on that read in any process that never called absl.app.run() – scripts, notebooks, pytest – so the data pipeline in docs/basics/training.md dies on its first batch. Marking the single flag as present makes absl return its default without touching the global “parsed” state; it is the same trick grain uses for its own debug flags. grain’s main branch already guards the read; this is a no-op once that ships, and a no-op if the flag is absent or flags were really parsed.

Return type:

None

gensbi.__version__