Colin Carroll

Possibly helpful as background is [this paper](https://arxiv.org/abs/2104.14421) (joint work with some of the researchers on the team), and a related [competition](https://izmailovpavel.github.io/neurips_bdl_competition/).

I think you have the batch dimension wrong for `cat` in the mixture: ```python zip = tfd.Mixture(cat = tfd.Categorical(probs=tf.stack([nonzero_prob, 1 - nonzero_prob], -1)), components=[tfd.Deterministic(tf.zeros_like(rate)), tfd.Poisson(rate)]) ``` at least has the...

I'll add that this adds a requirement on [dm-tree](https://github.com/google-deepmind/tree). I'd rather have used `jax.tree_utils`, but that requires all of `jax` (~1.4MB) instead of just `dm-tree` (~35kb).

It looks like maybe: https://github.com/google-deepmind/tree/issues/109 I gave https://github.com/google-deepmind/tree/pull/111 a try.

I tried it out, and it fails when trying to install `dm-tree` on python 3.12, because it requires `CMake`. Once you `conda install CMake` and install `xcode` on OSX, you...

Update: The PR was merged for `dm-tree`, but we need to wait on a pypi release.

Hey! They sneakily added PyPI wheels for 3.12 to dm-tree. I think this is now ready to go.

I think this is ready again.

Sounds good! `bayeux` will probably pin to the next release that has this in it, since [this example](https://github.com/jax-ml/bayeux/blob/main/docs/examples/oryx_and_bayeux.ipynb) produces a dictionary of parameters that the current `az.from_dict` fails to parse...

LGTM! On Wed, Mar 13, 2024, 3:16 PM Oriol Abril-Pla ***@***.***> wrote: > I think it should be done now. Docs used to look like this: > > Captura.de.pantalla.de.2024-03-13.19-42-15.png (view...