From f696ae5b9a18f0e1d625f34f6b9b408bfc9c1c31 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Thu, 17 Apr 2025 14:13:07 +0200
Subject: [PATCH] python312Packages.waymax: adapt to jax 0.6.0 new API changes

https://github.com/waymo-research/waymax/pull/77
---
 .../python-modules/waymax/default.nix           | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/pkgs/development/python-modules/waymax/default.nix b/pkgs/development/python-modules/waymax/default.nix
index 1e499b92f03a..9ece18dacaee 100644
--- a/pkgs/development/python-modules/waymax/default.nix
+++ b/pkgs/development/python-modules/waymax/default.nix
@@ -19,7 +19,7 @@
   tqdm,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage {
   pname = "waymax";
   version = "0-unstable-2025-03-25";
   pyproject = true;
@@ -31,6 +31,21 @@ buildPythonPackage rec {
     hash = "sha256-B1Rp5MATbEelp6G6K2wwV83QpINhOHgvAxb3mBN52Eg=";
   };
 
+  # AttributeError: jax.tree_map was removed in JAX v0.6.0: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version).
+  # https://github.com/waymo-research/waymax/pull/77
+  postPatch = ''
+    substituteInPlace \
+      waymax/agents/expert.py \
+      waymax/agents/waypoint_following_agent.py \
+      waymax/agents/waypoint_following_agent_test.py \
+      waymax/dynamics/abstract_dynamics_test.py \
+      waymax/dynamics/state_dynamics_test.py \
+      waymax/env/base_environment_test.py \
+      waymax/env/rollout_test.py \
+      waymax/env/wrappers/brax_wrapper_test.py \
+      --replace-fail "jax.tree_map" "jax.tree_util.tree_map"
+  '';
+
   build-system = [ setuptools ];
 
   dependencies = [