"""compat.py — brax 0.14 still calls jax.device_put_replicated, removed in jax 0.11.
Import this before brax. Replicates a pytree across devices with a leading device axis,
which is what brax's pmap-based training loop expects."""
import jax, jax.numpy as jnp, numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

if not hasattr(jax, "device_put_replicated"):
    def device_put_replicated(x, devices):
        devices = list(devices)
        mesh = Mesh(np.array(devices), ("d",))
        stacked = jax.tree_util.tree_map(lambda a: jnp.broadcast_to(jnp.asarray(a), (len(devices),) + jnp.shape(a)), x)
        return jax.device_put(stacked, NamedSharding(mesh, P("d")))
    jax.device_put_replicated = device_put_replicated