python312Packages.waymax: 0-unstable-2024-03-23 -> 0-unstable-2025-05-30 (#413545)

pkgs/development/python-modules/waymax/default.nix

@@ -21,31 +21,16 @@
 
 buildPythonPackage {
   pname = "waymax";
-  version = "0-unstable-2024-03-23";
+  version = "0-unstable-2025-05-30";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "waymo-research";
     repo = "waymax";
-    rev = "720f9214a9bf79b3da7926497f0cd0468ca3e630";
-    hash = "sha256-B1Rp5MATbEelp6G6K2wwV83QpINhOHgvAxb3mBN52Eg=";
+    rev = "48b33d71aac20a22db7d25f2d3220596899d944a";
+    hash = "sha256-YV0KI0UrFXO3HvKjqJE+K+hJJuYI4GiIR4l1fZNnl/E=";
   };
 
-  # 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 = [