python312Packages.waymax: adapt to jax 0.6.0 new API changes

https://github.com/waymo-research/waymax/pull/77

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 = [