bamojax 0.4.2__tar.gz

bamojax-0.4.2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Max Hinne, Uncertainty in complex systems
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

bamojax-0.4.2/MANIFEST.in

@@ -0,0 +1,4 @@
+include versioneer.py
+include README.md
+include LICENSE
+recursive-include bamojax *.py

bamojax-0.4.2/PKG-INFO

@@ -0,0 +1,264 @@
+Metadata-Version: 2.4
+Name: bamojax
+Version: 0.4.2
+Summary: Bayesian Modelling with JAX
+Author: Max Hinne
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/UncertaintyInComplexSystems/bamojax
+Project-URL: Issues, https://github.com/UncertaintyInComplexSystems/bamojax/issues
+Project-URL: Documentation, https://mhinne.github.io/uncertainty-in-complex-systems/bamojax_demo.html
+Keywords: bayesian modelling,jax,probabilistic,mcmc,smc,inference
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.10.0
+Requires-Dist: jaxlib>=0.10.0
+Requires-Dist: blackjax>=1.6
+Requires-Dist: numpyro>=0.22
+Requires-Dist: jaxtyping
+Requires-Dist: optax
+Requires-Dist: tqdm
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: pytest-cov>=5; extra == "dev"
+Requires-Dist: ipython>=8.0; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: mkdocs; extra == "docs"
+Requires-Dist: mkdocs-material; extra == "docs"
+Requires-Dist: mkdocstrings[python]; extra == "docs"
+Requires-Dist: pymdown-extensions; extra == "docs"
+Requires-Dist: mkdocs-gen-files; extra == "docs"
+Requires-Dist: mkdocs-literate-nav; extra == "docs"
+Dynamic: license-file
+
+[![Tests](https://github.com/uncertaintyincomplexsystems/bamojax/actions/workflows/test.yml/badge.svg)](https://github.com/uncertaintyincomplexsystems/bamojax/actions/workflows/test.yml) [![DOI](https://joss.theoj.org/papers/10.21105/joss.08642/status.svg)](https://doi.org/10.21105/joss.08642)
+
+![bamojax logo](https://github.com/uncertaintyincomplexsystems/bamojax/blob/main/figures/bamojax_logo_2x.png "bamojax")
+
+# Welcome to bamojax!
+
+Welcome to **bamojax**, the Bayesian modelling and model comparison toolbox implemented using the JAX coding universe. **bamojax** is a probabilistic programming language (PPL), similar to Numpyro, PyMC, Stan, JAGS, and BUGS, as well as a set of analyses for model comparison and averaging. **bamojax** relies on [Blackjax](https://blackjax-devs.github.io/blackjax/) for approximate inference, and on [Numpyro](https://num.pyro.ai/en/stable/) for probability distributions and their essential operations. 
+
+## What sets bamojax apart?
+
+Existing PPLs, such as PyMC, can export their log density function so it can be sampled using Blackjax. However, this has two downsides:
+
+1. Not all existing probabilistic programming easily allow you to update variables in your model using Gibbs Markov chain Monte Carlo (MCMC) algorithms. For example, if one wants to approximate the posterior over a latent Gaussian process (GP) and its hyperparameters, it can be more efficient to use elliptical slice sampling for the GP, than to apply the No-U-Turn Hamiltonian Monte Carlo Sampler [(Hoffman and Gelman, 2014)](https://jmlr.org/papers/v15/hoffman14a.html) to all variables at once. This effect is even more pronounced when embedding MCMC sampling within Sequential Monte Carlo (see [Hinne, 2025](https://link.springer.com/article/10.3758/s13428-025-02642-1) for more details).
+2. It is harder to embed this logdensity into tempered Sequential Monte Carlo algorithms, as for this one needs access to the prior and likelihood separately. 
+
+By implementing your own models and samplers using Blackjax, these problems can be circumvented. However, this is a labour-intensive and error-prone process. Therefore, **bamojax** provides a user-friendly interface around Blackjax, that allows for easy model construction and Gibbs sampling. 
+
+Once a Bayesian model has been determined, the next goal is often to compare it with competing hypotheses. This can, for example, be done by computing the marginal likelihoods of different models, and taking their ratio, a quantity known as the _Bayes factor_. However, computing marginal likelihoods and therefore Bayes factors is computationally challenging, and several approximation strategies exist. **bamojax** makes this process easier by implementing the following marginal likelihood approximations:
+
+- Importance sampling
+- Sequential Monte Carlo
+- The Laplace approximation
+- Bridge sampling (Gronau et al., 2017; Meng & Wong, 1996)
+- The THAMES estimator (Metodiev et al., 2025)
+
+More comparison features are being added to the toolbox continuously.
+
+## Installation
+
+Install **bamojax** using:
+
+```
+pip install git+https://github.com/UncertaintyInComplexSystems/bamojax#egg=bamojax
+```
+
+Please note that if you want to use hardware acceleration, such as using GPUs, you might want to install `JAX` yourself. Please see https://docs.jax.dev/en/latest/installation.html for instructions. If no JAX installation is present, **bamojax** will install it as a dependency.
+
+
+## Quick tutorial
+
+Let's estimate the latent probability $\theta=p(x_1=\text{heads} \mid \theta)$ that a coin lands up heads, given a set of observations $x_1, \ldots, x_n$:
+
+#### Generate data
+``` 
+import jax.numpy as jnp
+import jax.random as jrnd
+from bamojax.base import Model
+from bamojax.inference import SMCInference
+from bamojax.samplers import mcmc_sampler
+import numpyro as npr
+import numpyro.distributions as dist
+import blackjax as bjx
+
+key = jrnd.PRNGKey(0)
+key, key_data, key_inference = jrnd.split(key, 3)
+
+true_theta = 0.3
+n = 10000
+x = jrnd.bernoulli(key_data, p=true_theta, shape=(n, ))
+```
+
+#### Define Bayesian generative model
+
+Under the good, **bamojax** represents the a Bayesian model using a Directed Acyclic Graph (DAG) structure, and automatically derives model priors, likelihoods, and Gibbs inference schemes. These are then combined with the fast inference algorithms implemented in Blackjax. Bayesian models can be instantiated in **bamojax** using:
+
+```
+from bamojax.base import Model
+my_model = Model('The name of the model')
+``` 
+
+Subsequently, variables can be added to this model using for example:
+
+```
+def link_fn(probs):
+    return {'probs': probs}
+
+#
+latent_theta = my_model.add_node('theta', 
+                                 distribution=dist.Beta(concentration1=1, concentration0=1))
+observations = my_model.add_node('x', 
+                                 distribution=dist.Bernoulli, 
+                                 parents=dict(probs=latent_theta), 
+                                 link_fn=link_fn, 
+                                 observations=x)
+```
+
+#### Perform approximate inference
+
+The next step is to perform inference. Here, we use Adaptive-Tempered Sequential Monte Carlo, as implemented by Blackjax:
+
+```
+num_particles = 1_000
+num_mcmc_steps = 100
+num_chains = 1
+
+stepsize = 0.01
+
+mcmc_params = dict(sigma=stepsize*jnp.eye(my_model.get_model_size()))
+rmh = mcmc_sampler(my_model, 
+                   mcmc_kernel=bjx.normal_random_walk, 
+                   mcmc_parameters=mcmc_params)
+
+key, key_inference = jrnd.split(key)
+
+engine = SMCInference(model=my_model, 
+                      mcmc_kernel=rmh, 
+                      num_particles=num_particles, 
+                      num_mutations=num_mcmc_steps, 
+                      num_chains=num_chains)
+result = engine.run(key_inference)
+
+
+print(jnp.mean(result['final_state'].particles['theta']))            
+       
+>>> 0.29956531188263924
+```
+
+## Modelling in bamojax
+
+**bamojax** aims to make Bayesian modelling and inference with JAX and the Blackjax library more user-friendly, while retaining the speed and flexibility of the latter. To use the framework, a user is required to specify the directed acyclic graph that defines a Bayesian probabilistic model, as well as 
+* Distributions,
+* Transformations,
+* Link functions, and
+* Observations.
+
+### The model
+
+To perform these steps, one first initializes the `Model`:
+
+```
+my_model = Model('The name of my model')
+```
+
+This effectively creates an empty DAG. 
+
+### The variables
+
+Next, one adds nodes to the DAG one by one, or multiple at once in the case of iteration / plate notation: For instance, the following adds a latent random variable with a Gaussian distribution assigned to it, $x\sim \mathcal{N}(0,1)$:
+
+```
+a_new_node = my_model.add_node('variable name', distribution=dist.Normal, parents=dict(loc=0.0, scale=1.0))
+```
+
+Depending on whether a user provides values for the `distribution=`, and/or `observations=` arguments in the `add_node()` method, **bamojax** derives whether the variable is:
+1. Stochastic and latent,
+2. Stochastic and observed,
+3. Deterministic and observed.
+
+The `parents` argument expects a dictionary, in which the keys must correspond to either the arguments of a Numpyro `dist.Distribution` object or the arguments of a link function.
+
+Nodes can furthermore take any Numpyro bijector using the `bijector=` argument, for example to transform a real variable to a bounded domain or vice versa. In many cases, the same goal can be achieved by the link function, but sometimes a bijector is simpler to use.
+
+### Link functions
+
+An important 'feature' of **bamojax** is that it is straightforward to add any deterministic transformation from the values of a parent variable to the inputs of a child variable. For example, imagine a variable $\theta$ representing a coin flip probability. A typical prior would be the beta distribution, typically parametrized with pseudo-counts $\alpha$ and $\beta$, so that $\theta \sim \text{Beta}(\alpha,\beta)$.  However, we may want to specify a hierarchical prior on the _mode_ and _precision_ of this distribution, rather than on the pseudo-counts. With a link function we can express this:
+
+```
+def beta_link_fn(mode, conc):
+  a = mode*(conc-2) + 1
+  b = (1 - mode)*(conc-2) + 1
+  return {'concentration1': a, 'concentration0': b}
+
+omega = my_model.add_node('omega', distribution=dist.Beta, parents=dict(concentration1=1.0, concentration0=1.0))
+theta = my_model.add_node('theta', distribution=dist.Beta, parents=dict(mode=omega, conc=15), link_fn=beta_link_fn)
+```
+
+The link function `beta_link_fn` takes the mode (given by another node `omega`) and the concentration (given by a scalar, which is implicitly converted to a deterministic and observed node), and returns the standard arguments $\alpha$ and $\beta$ which the `dist.Beta` distribution object recognizes as valid parameters.
+
+Because link functions are written in plain Python, they can be of arbitrary complexity. For example, in [bamojax/examples/bnn/bnn_mpl.ipynb](https://github.com/UncertaintyInComplexSystems/bamojax/blob/main/bamojax/examples/bnn/bnn_mpl.ipynb), we use [Flax Linen](https://flax.readthedocs.io/en/latest/) to set up a multilayer perceptron as a link function.
+
+### Inference
+
+**bamojax** does not have its own inference engine, but provides an interface to [Blackjax](https://blackjax-devs.github.io/blackjax/), in addition to some quality-of-life features. Ultimately, control is left entirely to the user. Here is an example where we use Gibbs MCMC:
+
+```
+step_fns = dict(beta=normal_random_walk, sigma=normal_random_walk)
+step_fn_params = dict(omega=dict(sigma=0.05), theta=dict(sigma=0.5))
+gibbs_kernel = gibbs_sampler(my_model, step_fns=step_fns, step_fn_params=step_fn_params)
+```
+
+In Gibbs sampling, we specify a kernel for each variable. In this case, we use Gaussian proposal distributions for both $\omega$ and $\theta$.
+
+Now we can set up our inference engine:
+
+```
+engine = MCMCInference(model=my_model, num_chains=4, num_samples=100_000, num_burn=100_000, num_thin=50, mcmc_kernel=gibbs_kernel, return_diagnostics=True)
+result = engine.run(jrnd.PRNGKey(0))
+```
+
+This returns a dictionary with a `states` value containing a dictionary samples for each variable, for the requested number of chains and samples, discarding the specified number of burn-in samples, and storing only every 50th sample. By setting `return_diagnostics=True`, information such as acceptance rates are provided as well. For large models, turning this off can conserve memory consumption.
+
+Alternative inference engines include Sequential Monte Carlo, Variational Inference, and Stochastic-Gradient MCMC methods. Examples on how to use these can be found in the `examples/` folder, as well as on the [Uncertainty in Complex Systems website](https://mhinne.github.io/uncertainty-in-complex-systems).
+
+### Predictions
+
+**bamojax** supports each combination of sampling from the prior or posterior, and the latent variables or the predictive distribution, using any of the following:
+
+|                     |Prior                                         |Posterior                                                                         |
+|---------------------|----------------------------------------------|----------------------------------------------------------------------------------|
+|**Latent variables** |`my_model.sample_prior(key)`                  |Using `InferenceEngine`                                                           |
+|**Predictive**       |`my_model.sample_prior_predictive(key)`       |`my_model.sample_posterior_predictive(key, posterior_samples, input_variables)`   |
+
+When sampling from the posterior predictive, the parameter `input_variables=` can be used to provide for example predictor values, such as when sampling from a regression model or a neural network; $p(y^* \mid X, Y, x^*)$.
+
+## API
+
+The full **bamojax** API can be found here: https://uncertaintyincomplexsystems.github.io/bamojax/.
+
+## Citing bamojax
+
+If you use Bamojax in your work, please cite the following:
+
+    Max Hinne. (2025). Bamojax: Bayesian Modelling with JAX (Version 0.3.11). Journal of Open Source Software, 114. https://doi.org/10.21105/joss.08642.
+
+BibTeX:
+
+```
+@article{hinne2025bamojax,
+  author       = {Hinne, Max},
+  title        = {Bamojax: Bayesian Modelling with JAX},
+  journal      = {Journal of Open Source Software},
+  volume       = 114
+  year         = 2025,
+  publisher    = {Open Journals},
+  version      = {0.3.11},
+  doi          = {10.21105/joss.08642},
+  url          = {https://joss.theoj.org/papers/10.21105/joss.08642}
+}
+```

bamojax-0.4.2/README.md

@@ -0,0 +1,229 @@
+[![Tests](https://github.com/uncertaintyincomplexsystems/bamojax/actions/workflows/test.yml/badge.svg)](https://github.com/uncertaintyincomplexsystems/bamojax/actions/workflows/test.yml) [![DOI](https://joss.theoj.org/papers/10.21105/joss.08642/status.svg)](https://doi.org/10.21105/joss.08642)
+
+![bamojax logo](https://github.com/uncertaintyincomplexsystems/bamojax/blob/main/figures/bamojax_logo_2x.png "bamojax")
+
+# Welcome to bamojax!
+
+Welcome to **bamojax**, the Bayesian modelling and model comparison toolbox implemented using the JAX coding universe. **bamojax** is a probabilistic programming language (PPL), similar to Numpyro, PyMC, Stan, JAGS, and BUGS, as well as a set of analyses for model comparison and averaging. **bamojax** relies on [Blackjax](https://blackjax-devs.github.io/blackjax/) for approximate inference, and on [Numpyro](https://num.pyro.ai/en/stable/) for probability distributions and their essential operations. 
+
+## What sets bamojax apart?
+
+Existing PPLs, such as PyMC, can export their log density function so it can be sampled using Blackjax. However, this has two downsides:
+
+1. Not all existing probabilistic programming easily allow you to update variables in your model using Gibbs Markov chain Monte Carlo (MCMC) algorithms. For example, if one wants to approximate the posterior over a latent Gaussian process (GP) and its hyperparameters, it can be more efficient to use elliptical slice sampling for the GP, than to apply the No-U-Turn Hamiltonian Monte Carlo Sampler [(Hoffman and Gelman, 2014)](https://jmlr.org/papers/v15/hoffman14a.html) to all variables at once. This effect is even more pronounced when embedding MCMC sampling within Sequential Monte Carlo (see [Hinne, 2025](https://link.springer.com/article/10.3758/s13428-025-02642-1) for more details).
+2. It is harder to embed this logdensity into tempered Sequential Monte Carlo algorithms, as for this one needs access to the prior and likelihood separately. 
+
+By implementing your own models and samplers using Blackjax, these problems can be circumvented. However, this is a labour-intensive and error-prone process. Therefore, **bamojax** provides a user-friendly interface around Blackjax, that allows for easy model construction and Gibbs sampling. 
+
+Once a Bayesian model has been determined, the next goal is often to compare it with competing hypotheses. This can, for example, be done by computing the marginal likelihoods of different models, and taking their ratio, a quantity known as the _Bayes factor_. However, computing marginal likelihoods and therefore Bayes factors is computationally challenging, and several approximation strategies exist. **bamojax** makes this process easier by implementing the following marginal likelihood approximations:
+
+- Importance sampling
+- Sequential Monte Carlo
+- The Laplace approximation
+- Bridge sampling (Gronau et al., 2017; Meng & Wong, 1996)
+- The THAMES estimator (Metodiev et al., 2025)
+
+More comparison features are being added to the toolbox continuously.
+
+## Installation
+
+Install **bamojax** using:
+
+```
+pip install git+https://github.com/UncertaintyInComplexSystems/bamojax#egg=bamojax
+```
+
+Please note that if you want to use hardware acceleration, such as using GPUs, you might want to install `JAX` yourself. Please see https://docs.jax.dev/en/latest/installation.html for instructions. If no JAX installation is present, **bamojax** will install it as a dependency.
+
+
+## Quick tutorial
+
+Let's estimate the latent probability $\theta=p(x_1=\text{heads} \mid \theta)$ that a coin lands up heads, given a set of observations $x_1, \ldots, x_n$:
+
+#### Generate data
+``` 
+import jax.numpy as jnp
+import jax.random as jrnd
+from bamojax.base import Model
+from bamojax.inference import SMCInference
+from bamojax.samplers import mcmc_sampler
+import numpyro as npr
+import numpyro.distributions as dist
+import blackjax as bjx
+
+key = jrnd.PRNGKey(0)
+key, key_data, key_inference = jrnd.split(key, 3)
+
+true_theta = 0.3
+n = 10000
+x = jrnd.bernoulli(key_data, p=true_theta, shape=(n, ))
+```
+
+#### Define Bayesian generative model
+
+Under the good, **bamojax** represents the a Bayesian model using a Directed Acyclic Graph (DAG) structure, and automatically derives model priors, likelihoods, and Gibbs inference schemes. These are then combined with the fast inference algorithms implemented in Blackjax. Bayesian models can be instantiated in **bamojax** using:
+
+```
+from bamojax.base import Model
+my_model = Model('The name of the model')
+``` 
+
+Subsequently, variables can be added to this model using for example:
+
+```
+def link_fn(probs):
+    return {'probs': probs}
+
+#
+latent_theta = my_model.add_node('theta', 
+                                 distribution=dist.Beta(concentration1=1, concentration0=1))
+observations = my_model.add_node('x', 
+                                 distribution=dist.Bernoulli, 
+                                 parents=dict(probs=latent_theta), 
+                                 link_fn=link_fn, 
+                                 observations=x)
+```
+
+#### Perform approximate inference
+
+The next step is to perform inference. Here, we use Adaptive-Tempered Sequential Monte Carlo, as implemented by Blackjax:
+
+```
+num_particles = 1_000
+num_mcmc_steps = 100
+num_chains = 1
+
+stepsize = 0.01
+
+mcmc_params = dict(sigma=stepsize*jnp.eye(my_model.get_model_size()))
+rmh = mcmc_sampler(my_model, 
+                   mcmc_kernel=bjx.normal_random_walk, 
+                   mcmc_parameters=mcmc_params)
+
+key, key_inference = jrnd.split(key)
+
+engine = SMCInference(model=my_model, 
+                      mcmc_kernel=rmh, 
+                      num_particles=num_particles, 
+                      num_mutations=num_mcmc_steps, 
+                      num_chains=num_chains)
+result = engine.run(key_inference)
+
+
+print(jnp.mean(result['final_state'].particles['theta']))            
+       
+>>> 0.29956531188263924
+```
+
+## Modelling in bamojax
+
+**bamojax** aims to make Bayesian modelling and inference with JAX and the Blackjax library more user-friendly, while retaining the speed and flexibility of the latter. To use the framework, a user is required to specify the directed acyclic graph that defines a Bayesian probabilistic model, as well as 
+* Distributions,
+* Transformations,
+* Link functions, and
+* Observations.
+
+### The model
+
+To perform these steps, one first initializes the `Model`:
+
+```
+my_model = Model('The name of my model')
+```
+
+This effectively creates an empty DAG. 
+
+### The variables
+
+Next, one adds nodes to the DAG one by one, or multiple at once in the case of iteration / plate notation: For instance, the following adds a latent random variable with a Gaussian distribution assigned to it, $x\sim \mathcal{N}(0,1)$:
+
+```
+a_new_node = my_model.add_node('variable name', distribution=dist.Normal, parents=dict(loc=0.0, scale=1.0))
+```
+
+Depending on whether a user provides values for the `distribution=`, and/or `observations=` arguments in the `add_node()` method, **bamojax** derives whether the variable is:
+1. Stochastic and latent,
+2. Stochastic and observed,
+3. Deterministic and observed.
+
+The `parents` argument expects a dictionary, in which the keys must correspond to either the arguments of a Numpyro `dist.Distribution` object or the arguments of a link function.
+
+Nodes can furthermore take any Numpyro bijector using the `bijector=` argument, for example to transform a real variable to a bounded domain or vice versa. In many cases, the same goal can be achieved by the link function, but sometimes a bijector is simpler to use.
+
+### Link functions
+
+An important 'feature' of **bamojax** is that it is straightforward to add any deterministic transformation from the values of a parent variable to the inputs of a child variable. For example, imagine a variable $\theta$ representing a coin flip probability. A typical prior would be the beta distribution, typically parametrized with pseudo-counts $\alpha$ and $\beta$, so that $\theta \sim \text{Beta}(\alpha,\beta)$.  However, we may want to specify a hierarchical prior on the _mode_ and _precision_ of this distribution, rather than on the pseudo-counts. With a link function we can express this:
+
+```
+def beta_link_fn(mode, conc):
+  a = mode*(conc-2) + 1
+  b = (1 - mode)*(conc-2) + 1
+  return {'concentration1': a, 'concentration0': b}
+
+omega = my_model.add_node('omega', distribution=dist.Beta, parents=dict(concentration1=1.0, concentration0=1.0))
+theta = my_model.add_node('theta', distribution=dist.Beta, parents=dict(mode=omega, conc=15), link_fn=beta_link_fn)
+```
+
+The link function `beta_link_fn` takes the mode (given by another node `omega`) and the concentration (given by a scalar, which is implicitly converted to a deterministic and observed node), and returns the standard arguments $\alpha$ and $\beta$ which the `dist.Beta` distribution object recognizes as valid parameters.
+
+Because link functions are written in plain Python, they can be of arbitrary complexity. For example, in [bamojax/examples/bnn/bnn_mpl.ipynb](https://github.com/UncertaintyInComplexSystems/bamojax/blob/main/bamojax/examples/bnn/bnn_mpl.ipynb), we use [Flax Linen](https://flax.readthedocs.io/en/latest/) to set up a multilayer perceptron as a link function.
+
+### Inference
+
+**bamojax** does not have its own inference engine, but provides an interface to [Blackjax](https://blackjax-devs.github.io/blackjax/), in addition to some quality-of-life features. Ultimately, control is left entirely to the user. Here is an example where we use Gibbs MCMC:
+
+```
+step_fns = dict(beta=normal_random_walk, sigma=normal_random_walk)
+step_fn_params = dict(omega=dict(sigma=0.05), theta=dict(sigma=0.5))
+gibbs_kernel = gibbs_sampler(my_model, step_fns=step_fns, step_fn_params=step_fn_params)
+```
+
+In Gibbs sampling, we specify a kernel for each variable. In this case, we use Gaussian proposal distributions for both $\omega$ and $\theta$.
+
+Now we can set up our inference engine:
+
+```
+engine = MCMCInference(model=my_model, num_chains=4, num_samples=100_000, num_burn=100_000, num_thin=50, mcmc_kernel=gibbs_kernel, return_diagnostics=True)
+result = engine.run(jrnd.PRNGKey(0))
+```
+
+This returns a dictionary with a `states` value containing a dictionary samples for each variable, for the requested number of chains and samples, discarding the specified number of burn-in samples, and storing only every 50th sample. By setting `return_diagnostics=True`, information such as acceptance rates are provided as well. For large models, turning this off can conserve memory consumption.
+
+Alternative inference engines include Sequential Monte Carlo, Variational Inference, and Stochastic-Gradient MCMC methods. Examples on how to use these can be found in the `examples/` folder, as well as on the [Uncertainty in Complex Systems website](https://mhinne.github.io/uncertainty-in-complex-systems).
+
+### Predictions
+
+**bamojax** supports each combination of sampling from the prior or posterior, and the latent variables or the predictive distribution, using any of the following:
+
+|                     |Prior                                         |Posterior                                                                         |
+|---------------------|----------------------------------------------|----------------------------------------------------------------------------------|
+|**Latent variables** |`my_model.sample_prior(key)`                  |Using `InferenceEngine`                                                           |
+|**Predictive**       |`my_model.sample_prior_predictive(key)`       |`my_model.sample_posterior_predictive(key, posterior_samples, input_variables)`   |
+
+When sampling from the posterior predictive, the parameter `input_variables=` can be used to provide for example predictor values, such as when sampling from a regression model or a neural network; $p(y^* \mid X, Y, x^*)$.
+
+## API
+
+The full **bamojax** API can be found here: https://uncertaintyincomplexsystems.github.io/bamojax/.
+
+## Citing bamojax
+
+If you use Bamojax in your work, please cite the following:
+
+    Max Hinne. (2025). Bamojax: Bayesian Modelling with JAX (Version 0.3.11). Journal of Open Source Software, 114. https://doi.org/10.21105/joss.08642.
+
+BibTeX:
+
+```
+@article{hinne2025bamojax,
+  author       = {Hinne, Max},
+  title        = {Bamojax: Bayesian Modelling with JAX},
+  journal      = {Journal of Open Source Software},
+  volume       = 114
+  year         = 2025,
+  publisher    = {Open Journals},
+  version      = {0.3.11},
+  doi          = {10.21105/joss.08642},
+  url          = {https://joss.theoj.org/papers/10.21105/joss.08642}
+}
+```

bamojax-0.4.2/bamojax/__init__.py

@@ -0,0 +1,2 @@
+from . import _version
+__version__ = _version.get_versions()['version']

bamojax-0.4.2/bamojax/_version.py

@@ -0,0 +1,21 @@
+
+# This file was generated by 'versioneer.py' (0.29) from
+# revision-control system data, or from the parent directory name of an
+# unpacked source archive. Distribution tarballs contain a pre-generated copy
+# of this file.
+
+import json
+
+version_json = '''
+{
+ "date": "2026-10-26T21:35:56+0100",
+ "dirty": false,
+ "error": null,
+ "full-revisionid": "7b6576685f711fc94c568bb32813d45e67e2a861",
+ "version": "0.4.2"
+}
+'''  # END VERSION_JSON
+
+
+def get_versions():
+    return json.loads(version_json)