Spaces:

---

iteratehack

/

armins

Like 0

Sleeping

brax 0.14.2 (required by playground 0.2.0) still calls jax.device_put_replicated, which JAX deprecated in 0.8.1 and removed in 0.10.0. Only brax main has the sharding-based replacement.

Re-adds device_put_replicated and device_put_sharded when absent, using the same NamedSharding approach brax main uses. Chosen over pinning JAX below 0.10 because MJX/Warp already compile cleanly on the current stack and downgrading risks breaking that.

Also logs jax and brax versions at startup.