It's actually been super useful for us in several GP related models!

Bayes-by's first Markov chain? :o)

7/ Much more interesting examples, ranging from variable selection to Wishart processes, can be found at https://github.com/UncertaintyInComplexSystems/bamojax/tree/main/bamojax/examples

Happy modelling!

6/ model = Model('Gaussian with unknown mean')

unknown_mean = model.add_node('mu', distribution=dx.Normal(loc=mu0, scale=sd0))

y = model.add_node('y', distribution=dx.Normal, observations=y, parents=dict(loc=unknown_mean, scale=true_sd))

5/ Here is a very simple model, where we estimate the mean of a Gaussian using bamojax. We use Distrax to define probability distributions:

4/ In contrast to deriving a log-density from e.g. PyMC and using Blackjax to sample from it, with *bamojax* you have fine-grained control over individual Gibbs steps, which leads to much more efficient inference.

3/ This makes it possible to have *both* fast modelling, and fast inference!

2/ Very fast Bayesian inference in JAX is possible through Blackjax, but model development is easier in probabilistic programming languages like Stan and PyMC. *bamojax* provides automated Gibbs sampling, using the different sampling algorithms provided by Blackjax to update individual variables.