jax.device_put_replicated is deprecated, use jax.device_put instead

[zhuyeqing]

when i use brax 0.14 and jax 0.10, i will get error infomation like this : jax.device_put_replicated is deprecated, use jax.device_put instead

[RaspLeoPi]

downgrade jax to 0.9.2 for expedience