Hi,
I would like to know the expected release date for the next version of NumPyro. I have been using NumPyro extensively in my work, but encountered an issue with some production models when using the recent version of JAX.
It appears that this issue is addressed by #1890. However, we are restricted to using officially versioned libraries in our production environment.
Could you please provide any information regarding the release timeline?
We are currently blocked by Jitted function runs slower than non-jitted function for EM updates. · jax-ml/jax · Discussion #23822 · GitHub I think we can make a new release with XLA_FLAGS=--xla_cpu_use_thunk_runtime=false by default. Then remove it when the fix is available in jax.
Re timeline: probably at the end of the week
Thanks for the update. I’ve been also following the issue and currently using jax==0.4.31. Hope it gets resolved soon.