# Introducing bayeux

**URL:** <https://discourse.pymc.io/t/introducing-bayeux/13624>\
**Category:** Sharing\
**Created:** [January 12, 2024, 9:22pm UTC](https://discourse.pymc.io/t/introducing-bayeux/13624 "2024-01-12T21:22:20Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![colcarroll](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/colcarroll/32/9_2.png) [@colcarroll](https://discourse.pymc.io/u/colcarroll)\
**Post date:** [January 12, 2024, 9:22pm UTC](https://discourse.pymc.io/t/introducing-bayeux/13624/1 "2024-01-12T21:22:20Z")

</div>

Dear all – Happy to release a library for doing inference in JAX. It is ready to go as of ~20 minutes ago with (most) PyMC models. Check out [this quickstart colab](https://colab.research.google.com/drive/1a9sV0oJgb5_xH8N1y7MfRIFJ3sfKgfkP?usp=sharing) to jump right in (requires a recent version of PyMC), read [the docs here](https://jax-ml.github.io/bayeux/) or see [the code here](https://github.com/jax-ml/bayeux).

Currently it surfaces a bunch of ways of doing MCMC from blackjax and numpyro, as well as optimization from optax, jaxopt, and optimistix. There is also a VI routine from TFP that doesn’t quite work with all bayeux models. Coming up will be the VI routines from numpyro, and trying to work out why chees and meads from blackjax aren’t working very well on most models.

Happy to take pull requests, bug reports etc here or elsewhere.

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [January 12, 2024, 10:09pm UTC](https://discourse.pymc.io/t/introducing-bayeux/13624/2 "2024-01-12T22:09:18Z")

</div>

Feature request: Using with PyMC example

---

<div class="post-metadata">

**Author:** ![colcarroll](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/colcarroll/32/9_2.png) [@colcarroll](https://discourse.pymc.io/u/colcarroll)\
**Post date:** [January 12, 2024, 10:44pm UTC](https://discourse.pymc.io/t/introducing-bayeux/13624/3 "2024-01-12T22:44:25Z")

</div>

Yes! Tuesday morning something similar to the above colab will go up!

---

<div class="post-metadata">

**Author:** ![colcarroll](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/colcarroll/32/9_2.png) [@colcarroll](https://discourse.pymc.io/u/colcarroll)\
**Post date:** [January 17, 2024, 9:38pm UTC](https://discourse.pymc.io/t/introducing-bayeux/13624/4 "2024-01-17T21:38:40Z")

</div>

Update: It went up Wednesday morning instead! Adding this turned up a bug in `bayeux`’s structural vi from TFP implementation, but it _should_ now work generally.

> **[Using with PyMC - bayeux](https://jax-ml.github.io/bayeux/examples/pymc_and_bayeux/)**
>
> The documentation for the bayeux software library.

Next on the list is trying to add some VI from numpyro.

Probably not (yet) a good idea, but you could delete lots of `pymc/sampling/jax.py` and replace it with `bayeux` now 😁 (just don’t delete the `jaxify` functions that bayeux uses…)

---

<div class="post-metadata">

**Author:** ![zweli](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/zweli/32/4089_2.png) [@zweli](https://discourse.pymc.io/u/zweli)\
**Post date:** [January 28, 2024, 9:46pm UTC](https://discourse.pymc.io/t/introducing-bayeux/13624/5 "2024-01-28T21:46:07Z")

</div>

If you want to use this with Bambi just call model.build() and then use model.backend.model in the bx.Model.from\_pymc call like this:

```python
dist = pm.Normal.dist(mu=100, sigma=30)

draws = pm.draw(dist, draws=1000, random_seed=1000)

df = pd.DataFrame(data=draws, columns=['heights'])

formula = bmb.Formula('heights ~ 1')

model = bmb.Model(formula=formula, family='gaussian', data=df)

model.build()

bx_model = bx.Model.from_pymc(model.backend.model)

idata = bx_model.mcmc.numpyro_nuts(seed=jax.random.key(0))

az.summary(idata)

```

---

<div class="post-metadata">

**Author:** ![colcarroll](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/colcarroll/32/9_2.png) [@colcarroll](https://discourse.pymc.io/u/colcarroll)\
**Post date:** [January 29, 2024, 2:50am UTC](https://discourse.pymc.io/t/introducing-bayeux/13624/6 "2024-01-29T02:50:35Z")

</div>

Thank you! Would it be ok to add this to the documentation?

---

<div class="post-metadata">

**Author:** ![zweli](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/zweli/32/4089_2.png) [@zweli](https://discourse.pymc.io/u/zweli)\
**Post date:** [January 29, 2024, 3:23am UTC](https://discourse.pymc.io/t/introducing-bayeux/13624/7 "2024-01-29T03:23:40Z")

</div>

Of course!

---

<div class="post-metadata">

**Author:** ![joshualeond](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/joshualeond/32/8031_2.png) [@joshualeond](https://discourse.pymc.io/u/joshualeond)\
**Post date:** [April 10, 2024, 4:22am UTC](https://discourse.pymc.io/t/introducing-bayeux/13624/8 "2024-04-10T04:22:17Z")

</div>

Hi Colin! would this be the best path forward for using JAX + PyMC now?

I also see the PyMC docs [Faster Sampling with JAX and Numba](https://www.pymc.io/projects/examples/en/latest/samplers/fast_sampling_with_jax_and_numba.html) but that appears to be from 2022.

I’m a new PyMC user so please forgive me if this is obviously stated somewhere.

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [April 10, 2024, 5:52am UTC](https://discourse.pymc.io/t/introducing-bayeux/13624/9 "2024-04-10T05:52:48Z")

</div>

> [@joshualeond](#):
>
> I also see the PyMC docs [Faster Sampling with JAX and Numba](https://www.pymc.io/projects/examples/en/latest/samplers/fast_sampling_with_jax_and_numba.html) but that appears to be from 2022.

For most users that’s enough. bayeux gives you access to more wild/experimental samplers and optimizers but those documented in the docs suffice to 99% of the users.

Nothing wrong with exploring bayeux

---

<div class="post-metadata">

**Author:** ![colcarroll](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/colcarroll/32/9_2.png) [@colcarroll](https://discourse.pymc.io/u/colcarroll)\
**Post date:** [April 10, 2024, 1:08pm UTC](https://discourse.pymc.io/t/introducing-bayeux/13624/10 "2024-04-10T13:08:19Z")

</div>

Strong agree with @ricardoV94 – if you’re getting started with PyMC, you’ll find the best support, most responsive issue response, fastest bug fixes, and a huge amount of existing utility functions by just using the `nuts_sampler="numpyro"`. Once you call `bx.Model.from_pymc(model)`, your work downstream from there will be using the JAX ecosystem.

It suffices to be aware that you can usually add one line from `bayeux` and have access to some different samplers (and optimizers and VI).