I am just writing for a quick doubt . Flax is currently on pre-release state but hopefully soon to be fully released. Jax is on version 1.41 currently, (and Flax is only properly working with the latest version of Jax) while Numpyro right now is working with 1.37 version. Do you have an estimate of when everything could be working together, updated? Just wondering , I am aware is a lot of work.
Thanks and take care!