Rémi Louf

It's all in there: https://jax.readthedocs.io/en/latest/jax.experimental.host_callback.html#using-call-to-call-a-tensorflow-function-with-reverse-mode-autodiff-support

See this comment https://github.com/sphinx-doc/sphinx/issues/1514#issuecomment-742703082 in a Sphinx issue about how to achieve this with CSS.

See the following discussion https://github.com/blackjax-devs/blackjax/discussions/275#discussioncomment-3665608 for a template.

Rewrite the MNIST MLP example with Flax

3

https://blackjax-devs.github.io/blackjax/examples/SGMCMC.html And add SgHMC back!

Use `flax` for the Hierarchical BNN example

2

This should simplify the example code.

Refactor the HierarchicalBNN example

1

- Use `flax` to build the neural network; - The function `get_predictions` can be optimized, this may substantially speed up the notebook

Add several design documents so users can *understand* Blackjax

1

- [ ] Everything is a kernel - [ ] The 3-layer design - [ ] Programmable inference