run code with new version of brax and jax ?

[ucacaxm]

Hi,

Do you think the code may run with new version of brax and jax ?

I have: brax 0.10.5 brax-jumpy 0.1.0 flax 0.8.4 jax 0.4.29 jax-cuda12-pjrt 0.4.29 jax-cuda12-plugin 0.4.29 jaxlib 0.4.29 jaxopt 0.8.3 optax 0.2.2 orbax-checkpoint 0.5.16

I tested but I have issue with jumpy. I can search but I'd like to know if it's worth it or if too much has changed in Brax and Jax ?

Thanks