From 1aa137ea39a2bbf5da14cc68a9ec96ca3b2251af Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Mon, 20 Apr 2026 19:30:53 +0000
Subject: [PATCH] python3Packages.brax: fix jax 0.10.0 compatibility

---
 .../python-modules/brax/default.nix           | 15 ++++
 ...ont-use-device_put_replicated-compat.patch | 68 +++++++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100644 pkgs/development/python-modules/brax/dont-use-device_put_replicated-compat.patch

diff --git a/pkgs/development/python-modules/brax/default.nix b/pkgs/development/python-modules/brax/default.nix
index e2609c149536..d74a839c505a 100644
--- a/pkgs/development/python-modules/brax/default.nix
+++ b/pkgs/development/python-modules/brax/default.nix
@@ -40,6 +40,7 @@ buildPythonPackage (finalAttrs: {
   pname = "brax";
   version = "0.14.2";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "google";
@@ -48,6 +49,20 @@ buildPythonPackage (finalAttrs: {
     hash = "sha256-/oznBa44xKl+9T1YrTVhCbuKZj16V1BTlnmCGRbF45g=";
   };
 
+  patches = [
+    # AttributeError: jax.device_put_replicated is deprecated; use jax.device_put instead.
+    # See https://docs.jax.dev/en/latest/migrate_pmap.html#drop-in-replacements for a drop-in replacement.
+    ./dont-use-device_put_replicated-compat.patch
+  ];
+
+  # TypeError: clip() got an unexpected keyword argument 'a_min'
+  postPatch = ''
+    substituteInPlace brax/fluid.py \
+      --replace-fail \
+        "box = 6.0 * jp.clip(jp.sum(diag_inertia_v, axis=-1), a_min=1e-12)" \
+        "box = 6.0 * jp.clip(jp.sum(diag_inertia_v, axis=-1), min=1e-12)"
+  '';
+
   build-system = [
     hatchling
   ];
diff --git a/pkgs/development/python-modules/brax/dont-use-device_put_replicated-compat.patch b/pkgs/development/python-modules/brax/dont-use-device_put_replicated-compat.patch
new file mode 100644
index 000000000000..261aa5907a60
--- /dev/null
+++ b/pkgs/development/python-modules/brax/dont-use-device_put_replicated-compat.patch
@@ -0,0 +1,68 @@
+diff --git a/brax/training/agents/apg/train.py b/brax/training/agents/apg/train.py
+index f5fcb0e..87b198f 100644
+--- a/brax/training/agents/apg/train.py
++++ b/brax/training/agents/apg/train.py
+@@ -310,7 +310,7 @@ def train(
+           specs.Array((env.observation_size,), jnp.dtype(dtype))
+       ),
+   )
+-  training_state = jax.device_put_replicated(
++  training_state = pmap.device_put_replicated(
+       training_state, jax.local_devices()[:local_devices_to_use]
+   )
+ 
+diff --git a/brax/training/agents/ppo/train.py b/brax/training/agents/ppo/train.py
+index 9aec960..6624733 100644
+--- a/brax/training/agents/ppo/train.py
++++ b/brax/training/agents/ppo/train.py
+@@ -753,7 +753,7 @@ def train(
+         {},
+     )
+ 
+-  training_state = jax.device_put_replicated(
++  training_state = pmap.device_put_replicated(
+       training_state, jax.local_devices()[:local_devices_to_use]
+   )
+ 
+diff --git a/brax/training/agents/sac/train.py b/brax/training/agents/sac/train.py
+index be716e9..8dcf3bf 100644
+--- a/brax/training/agents/sac/train.py
++++ b/brax/training/agents/sac/train.py
+@@ -108,7 +108,7 @@ def _init_training_state(
+       alpha_params=log_alpha,
+       normalizer_params=normalizer_params,
+   )
+-  return jax.device_put_replicated(
++  return pmap.device_put_replicated(
+       training_state, jax.local_devices()[:local_devices_to_use]
+   )
+ 
+diff --git a/brax/training/pmap.py b/brax/training/pmap.py
+index 82760fc..af62ef8 100644
+--- a/brax/training/pmap.py
++++ b/brax/training/pmap.py
+@@ -19,12 +19,23 @@ from typing import Any
+ 
+ import jax
+ import jax.numpy as jnp
++from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
++import numpy as np
++
++
++def device_put_replicated(x, devices):
++  """Drop-in replacement for jax.device_put_replicated supporting pytrees."""
++  mesh = Mesh(np.array(devices), ('x',))
++  sharding = NamedSharding(mesh, P('x'))
++  return jax.tree.map(
++      lambda y: jax.device_put(jnp.stack([y] * len(devices)), sharding), x
++  )
+ 
+ 
+ def bcast_local_devices(value, local_devices_to_use=1):
+   """Broadcasts an object to all local devices."""
+   devices = jax.local_devices()[:local_devices_to_use]
+-  return jax.device_put_replicated(value, devices)
++  return device_put_replicated(value, devices)
+ 
+ 
+ def synchronize_hosts():