👋 Welcome to sbijax!

Simulation-based inference in JAX

---

sbijax implements several algorithms for simulation-based inference in JAX using Haiku, Distrax and BlackJAX. Specifically, sbijax implements

• Sequential Monte Carlo ABC (SMCABC)

• Neural Likelihood Estimation (SNL)

• Surjective Neural Likelihood Estimation (SSNL)

• Neural Posterior Estimation C (short SNP)

• Contrastive Neural Ratio Estimation (short SNR)

• Neural Approximate Sufficient Statistics (SNASS)

• Neural Approximate Slice Sufficient Statistics (SNASSS)

• Flow matching posterior estimation (SFMPE)

• Consistency model posterior estimation (SCMPE)

Caution

⚠️ As per the LICENSE file, there is no warranty whatsoever for this free software tool. If you discover bugs, please report them.

Installation

To install from PyPI, call:

pip install sbijax

To install the latest GitHub <RELEASE>, just call the following on the command line:

pip install git+https://github.com/dirmeier/sbijax@<RELEASE>

See also the installation instructions for JAX, if you plan to use sbijax on GPU/TPU.

Contributing

Contributions in the form of pull requests are more than welcome. A good way to start is to check out issues labelled “good first issue”.

In order to contribute:

1. Clone sbijax and install hatch via pip install hatch,

2. create a new branch locally git checkout -b feature/my-new-feature or git checkout -b issue/fixes-bug,

3. implement your contribution and ideally a test case,

4. test it by calling hatch run test on the (Unix) command line,

5. submit a PR 🙂

Acknowledgements

Note

📝 The API of the package is heavily inspired by the excellent Pytorch-based sbi package which is substantially more feature-complete and user-friendly.

License

sbijax is licensed under the Apache 2.0 License.