Fix: shim jax.device_put_replicated removed in JAX 0.10

app.py

@@ -53,6 +53,50 @@ def _upload(path: Path) -> None:
         log(f"upload failed ({type(e).__name__}): {e}")
 
 
+def install_jax_pmap_shims() -> None:
+    """Re-add jax.device_put_replicated / device_put_sharded if this JAX removed them.
+
+    JAX deprecated both in 0.8.1 and removed them in 0.10.0 (April 2026), but the
+    current brax *release* (0.14.2, which playground 0.2.0 requires) still calls
+    device_put_replicated; only brax main has the sharding-based replacement.
+    Shimming the two functions is a smaller intervention than pinning JAX down,
+    which would risk the MJX/Warp stack that already compiles cleanly here.
+    """
+    import jax
+    import jax.numpy as jnp
+    import numpy as np
+
+    def _sharding(devices):
+        mesh = jax.sharding.Mesh(np.array(list(devices)), axis_names=("i",))
+        return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("i"))
+
+    def device_put_replicated(x, devices):
+        sharding, n = _sharding(devices), len(devices)
+
+        def rep(leaf):
+            stack = jnp.stack if isinstance(leaf, jax.Array) else np.stack
+            return jax.device_put(stack([leaf] * n), sharding)
+
+        return jax.tree_util.tree_map(rep, x)
+
+    def device_put_sharded(shards, devices):
+        sharding = _sharding(devices)
+
+        def put(*leaves):
+            stack = jnp.stack if isinstance(leaves[0], jax.Array) else np.stack
+            return jax.device_put(stack(list(leaves)), sharding)
+
+        return jax.tree_util.tree_map(put, *shards)
+
+    for name, fn in (("device_put_replicated", device_put_replicated),
+                     ("device_put_sharded", device_put_sharded)):
+        try:
+            getattr(jax, name)
+        except AttributeError:
+            setattr(jax, name, fn)
+            log(f"shimmed jax.{name} (removed in this JAX version)")
+
+
 def train_worker() -> None:
     try:
         STATE["status"] = "importing"
@@ -60,7 +104,10 @@ def train_worker() -> None:
         import functools
         import jax
 
-        log(f"jax devices: {jax.devices()}")
+        log(f"jax {jax.__version__}  devices: {jax.devices()}")
+        install_jax_pmap_shims()
+        import brax
+        log(f"brax {brax.__version__}")
         if not any(d.platform == "gpu" for d in jax.devices()):
             log("WARNING: no GPU visible to JAX -- this will be very slow")