diff --git a/pkgs/development/python-modules/distrax/default.nix b/pkgs/development/python-modules/distrax/default.nix index 6c9d585a09f5..014167bf2aea 100644 --- a/pkgs/development/python-modules/distrax/default.nix +++ b/pkgs/development/python-modules/distrax/default.nix @@ -2,6 +2,7 @@ lib, buildPythonPackage, fetchFromGitHub, + fetchpatch, chex, jaxlib, numpy, @@ -23,6 +24,15 @@ buildPythonPackage rec { hash = "sha256-A1aCL/I89Blg9sNmIWQru4QJteUTN6+bhgrEJPmCrM0="; }; + patches = [ + # TODO: remove at the next release (already on master) + (fetchpatch { + name = "fix-jax-0.6.0-compat"; + url = "https://github.com/google-deepmind/distrax/commit/c02708ac46518fac00ab2945311e0f2ee32c672c.patch"; + hash = "sha256-hFNXKoA1b5I6dzhwTRXp/SnkHv89GI6tYwlnBBHwG78="; + }) + ]; + dependencies = [ chex jaxlib @@ -71,6 +81,10 @@ buildPythonPackage rec { ]; disabledTestPaths = [ + # Since jax 0.6.0: + # TypeError: () got an unexpected keyword argument 'accuracy' + "distrax/_src/bijectors/lambda_bijector_test.py" + # TypeErrors "distrax/_src/bijectors/tfp_compatible_bijector_test.py" "distrax/_src/distributions/distribution_from_tfp_test.py" diff --git a/pkgs/development/python-modules/dm-haiku/default.nix b/pkgs/development/python-modules/dm-haiku/default.nix index 5ea929f8570e..0a818b0b8f36 100644 --- a/pkgs/development/python-modules/dm-haiku/default.nix +++ b/pkgs/development/python-modules/dm-haiku/default.nix @@ -58,6 +58,17 @@ let }) ]; + # AttributeError: jax.core.Var was removed in JAX v0.6.0. Use jax.extend.core.Var instead, and + # see https://docs.jax.dev/en/latest/jax.extend.html for details. + # Alrady on master: https://github.com/google-deepmind/dm-haiku/commit/cfe8480d253a93100bf5e2d24c40435a95399c96 + # TODO: remove at the next release + postPatch = '' + substituteInPlace haiku/_src/jaxpr_info.py \ + --replace-fail "jax.core.JaxprEqn" "jax.extend.core.JaxprEqn" \ + --replace-fail "jax.core.Var" "jax.extend.core.Var" \ + --replace-fail "jax.core.Jaxpr" "jax.extend.core.Jaxpr" + ''; + build-system = [ setuptools ]; dependencies = [ diff --git a/pkgs/development/python-modules/gymnasium/default.nix b/pkgs/development/python-modules/gymnasium/default.nix index a87833f8bb63..c98c3d7cb119 100644 --- a/pkgs/development/python-modules/gymnasium/default.nix +++ b/pkgs/development/python-modules/gymnasium/default.nix @@ -90,6 +90,15 @@ buildPythonPackage rec { ]; disabledTests = [ + # Fails since jax 0.6.0 + # Fixed on master https://github.com/Farama-Foundation/Gymnasium/commit/94019feee1a0f945b9569cddf62780f4e1a224a5 + # TODO: un-skip at the next release + "test_all_env_api" + "test_env_determinism_rollout" + "test_jax_to_numpy_wrapper" + "test_pickle_env" + "test_roundtripping" + # Succeeds for most environments but `test_render_modes[Reacher-v4]` fails because it requires # OpenGL access which is not possible inside the sandbox. "test_render_mode" diff --git a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix index 14c59787f06a..160158d0ed39 100644 --- a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix @@ -2,7 +2,7 @@ lib, stdenv, buildPythonPackage, - fetchurl, + fetchPypi, addDriverRunpath, autoPatchelfHook, pypaInstallHook, @@ -31,30 +31,31 @@ let ] ); - # Find new releases at https://storage.googleapis.com/jax-releases - # When upgrading, you can get these hashes from jaxlib/prefetch.sh. See - # https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index. - - # upstream does not distribute jax-cuda12-pjrt binaries for aarch64-linux - srcs = { - "x86_64-linux" = fetchurl { - url = "https://storage.googleapis.com/jax-releases/cuda12_plugin/jax_cuda12_pjrt-${version}-py3-none-manylinux2014_x86_64.whl"; - hash = "sha256-xTeDBlaLoMgbIwp3ndMZTJ3RAzmrY2CugJKBCNN+f3U="; - }; - # "aarch64-linux" = fetchurl { - # url = "https://storage.googleapis.com/jax-releases/cuda12_plugin/jax_cuda12_pjrt-${version}-py3-none-manylinux2014_aarch64.whl"; - # hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA="; - # }; - }; in -buildPythonPackage { +buildPythonPackage rec { pname = "jax-cuda12-pjrt"; inherit version; pyproject = false; - src = - srcs.${stdenv.hostPlatform.system} - or (throw "jax-cuda12-pjrt: No src for ${stdenv.hostPlatform.system}"); + src = fetchPypi { + pname = "jax_cuda12_pjrt"; + inherit version; + format = "wheel"; + python = "py3"; + dist = "py3"; + platform = + { + x86_64-linux = "manylinux2014_x86_64"; + aarch64-linux = "manylinux2014_aarch64"; + } + .${stdenv.hostPlatform.system}; + hash = + { + x86_64-linux = "sha256-aDcb2cE1JEuJZjA5viCCVWmKdb7JhU1BnqPD+VfKRkY= "; + aarch64-linux = "sha256-m/67BqOWFMtomfdzDqhWHxEVasgcuz7GiEpir7OxX/M="; + } + .${stdenv.hostPlatform.system}; + }; nativeBuildInputs = [ autoPatchelfHook @@ -97,7 +98,7 @@ buildPythonPackage { sourceProvenance = [ lib.sourceTypes.binaryNativeCode ]; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ natsukium ]; - platforms = lib.attrNames srcs; + platforms = lib.platforms.linux; # see CUDA compatibility matrix # https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder broken = diff --git a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix index a1bcf01218da..3103feb271e4 100644 --- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix @@ -40,42 +40,42 @@ let "3.10-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp310"; - hash = "sha256-uiVVln+bbDgci075+wPQW8Vewl7P7lz+RcWs4099QVI="; + hash = "sha256-pwDhcYI84lUQIALkDJR4j6ho8hYle30/BWjQn+dcEHs="; }; "3.10-aarch64-linux" = getSrcFromPypi { platform = "manylinux2014_aarch64"; dist = "cp310"; - hash = "sha256-YXGu0vSzvdX8E3gt4QcsamNPzhNzG3XQywpquPTm5lA="; + hash = "sha256-UwrYUcpGKZHOgtsmrUfwKwjOvkg8nI0MADfp4np7Up8="; }; "3.11-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp311"; - hash = "sha256-qqcEpe9UdZXQItscHkh4oGdxFkEqk2DBFdZ/9LZOFZY="; + hash = "sha256-DZ7O3mbEAlhwKkImHoaM21ahA1UafDyISzX1Mcms1I4="; }; "3.11-aarch64-linux" = getSrcFromPypi { platform = "manylinux2014_aarch64"; dist = "cp311"; - hash = "sha256-KY0tdo8QKbdKCx0BJw5Uk0nSw33Adlh5ZULNqWfre9M="; + hash = "sha256-fNG0iKVKMInolYjMr2dwiZUsglKefQQD4LBQGZ5SVBg="; }; "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp312"; - hash = "sha256-IDDPEgjOTqcO5WysYd3SOfl5hpX8Obt3OcUKJdbp2kQ="; + hash = "sha256-5w608IRpbD474StekJ7xIFyfVu/j3OzyYhvZtatZVNU="; }; "3.12-aarch64-linux" = getSrcFromPypi { platform = "manylinux2014_aarch64"; dist = "cp312"; - hash = "sha256-wlF6fCGG+HCIlGluJs+W69YLeHnOyjmLLEarso0slsg="; + hash = "sha256-oqOvX5iIDYb40kartGpVLlou9J12e/xKdMjDV3UgB8Y="; }; "3.13-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp313"; - hash = "sha256-GGJZWyttgVZ50R4OiJ5SMYXuVKRtRuAiaJ9w/EVU3ZE="; + hash = "sha256-6W891KlCUWroeMn2l+au/teOFI8JAYynPuKLI0JqfYo="; }; "3.13-aarch64-linux" = getSrcFromPypi { platform = "manylinux2014_aarch64"; dist = "cp313"; - hash = "sha256-If7BtWyYeD6gVpt0elZ1Hx+f8hh7SKzBHHANO/xeGjE="; + hash = "sha256-o0LyznxLH1nUA/Zlo1qGuGUCU7sl3jRkf7IlxFzrCgQ="; }; }; in diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index 30b6789692a9..b8168916df8d 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -40,7 +40,7 @@ let in buildPythonPackage rec { pname = "jax"; - version = "0.5.3"; + version = "0.6.0"; pyproject = true; src = fetchFromGitHub { @@ -48,7 +48,7 @@ buildPythonPackage rec { repo = "jax"; # google/jax contains tags for jax and jaxlib. Only use jax tags! tag = "jax-v${version}"; - hash = "sha256-t4LHwpCz08zrQGWBehyPs2JnxsOvtV3L14MCdTqMeEI="; + hash = "sha256-leadxLK21JF/cF/hMveUPgrwyumjR+zMm9Wsn7bWNLQ="; }; build-system = [ setuptools ]; diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix index 4a8bfb8d31af..8c934c4af977 100644 --- a/pkgs/development/python-modules/jaxlib/bin.nix +++ b/pkgs/development/python-modules/jaxlib/bin.nix @@ -18,7 +18,7 @@ }: let - version = "0.5.3"; + version = "0.6.0"; inherit (python) pythonVersion; # As of 2023-06-06, google/jax upstream is no longer publishing CPU-only wheels to their GCS bucket. Instead the @@ -49,65 +49,65 @@ let "3.10-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp310"; - hash = "sha256-Ur5sl3Wv9zimEXDYwEdQXHW7eZpFUY5mp6CQgSexF4U="; + hash = "sha256-pNQlTHEziIh6MhN508WxogITqNzckD+vFRObqB4+zWE="; }; "3.10-aarch64-linux" = getSrcFromPypi { platform = "manylinux2014_aarch64"; dist = "cp310"; - hash = "sha256-lyQA20r26FJw2B215uYg0xOV8EcuUQxQ381Ms/crciA="; + hash = "sha256-VBpBi5iyjfW9Oh6TxistP2TUSwxwt7YI9/4rSqRSsq8="; }; "3.10-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp310"; - hash = "sha256-SP9cifuKD+BNR16d3AdLSHmpHXq2ilHOxc0eh/gebEc="; + hash = "sha256-ZKgvjrQP23uh1G75BzANQuT5jL2pYCou2OcNsamsSmA="; }; "3.11-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp311"; - hash = "sha256-KeFTD8gYMyFvHii1eNDFlpdlT3LuMcekTtd1O69axGY="; + hash = "sha256-vtRVJeO7XsCGML/SB8Ca+dYun/E/XwfC7iz9jthBG6E="; }; "3.11-aarch64-linux" = getSrcFromPypi { platform = "manylinux2014_aarch64"; dist = "cp311"; - hash = "sha256-pGZvgdcsBg7T5YHe0RapyqmwpwoUilTLEqHTr8o2JLU="; + hash = "sha256-wK6VmJmALhMpzI7ForTUvpoHa1vrIFLrSbo3UU5iPrw="; }; "3.11-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp311"; - hash = "sha256-tivYsp5aT5v6pXyNr24EggssmU9Ejz3sYC1kJVVF6fI="; + hash = "sha256-7xY88H3gC8VpAWnpf6+q3DePHDgfAofopHPnirW6sbU="; }; "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp312"; - hash = "sha256-Wl6IqxzW/feNaavjVE6PCcziAN0zm7hfvjwupn8qXmg="; + hash = "sha256-tthbjR/XkkiwRQNRcgHnL8vNOYDPeR036BRwnqUKPII="; }; "3.12-aarch64-linux" = getSrcFromPypi { platform = "manylinux2014_aarch64"; dist = "cp312"; - hash = "sha256-vd9jYDd6oceS5H/YfzB8NC4zHl/zWC+UCxvKAPa0vHM="; + hash = "sha256-JTb6k+wUjVAW2osgd7pmMlsNhqriKJphwSaHfwQrPRw="; }; "3.12-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp312"; - hash = "sha256-05Tb3koca9Z1Ac+ynTgZoQuQDLU0zA/GAzGfcJLyTPo="; + hash = "sha256-fjzi7w7cm0izbicEw2GB8eznoSrBFN91PbQobqLG6Lg="; }; "3.13-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp313"; - hash = "sha256-u3WTy3//yxOWPyL6UintlguPtK5ew7CCAEjL1n8ejjE="; + hash = "sha256-0PsSLceDDKKlyjyHSghzY6AFMrZEUJwhnDv9HVRRXo0="; }; "3.13-aarch64-linux" = getSrcFromPypi { platform = "manylinux2014_aarch64"; dist = "cp313"; - hash = "sha256-6QS5Le37x+VFclqNdnaYcDCunAaQAdlHAbwQnG2rQQA="; + hash = "sha256-GJcpY5diBQwXgLBQ6Y/2IEgLHqMr8WdTPgAKXPTFc44="; }; "3.13-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp313"; - hash = "sha256-MTIcJSgqBqbfyUBQe8FNCgrIONjO1sB6oAp/rjTOez8="; + hash = "sha256-xOl5NMuvUXI0OqWujvDFhGLOJhVN/adUICswNBYMrHs="; }; }; in diff --git a/pkgs/development/python-modules/rlax/default.nix b/pkgs/development/python-modules/rlax/default.nix index c74c35c76ed9..7d800cc5e62f 100644 --- a/pkgs/development/python-modules/rlax/default.nix +++ b/pkgs/development/python-modules/rlax/default.nix @@ -36,6 +36,7 @@ buildPythonPackage rec { hash = "sha256-v2Lbzya+E9d7tlUVlQQa4fuPp2q3E309Qvyt70mcdb0="; }; + # TODO: remove these patches at the next release (already on master) patches = [ (fetchpatch { # Follow chex API change (https://github.com/google-deepmind/chex/pull/52) @@ -43,6 +44,22 @@ buildPythonPackage rec { url = "https://github.com/google-deepmind/rlax/commit/30e7913a1102667137654d6e652a6c4b9e9ba1f4.patch"; hash = "sha256-OPnuTKEtwZ28hzR1660v3DcktxTYjhR1xYvFbQvOhgs="; }) + (fetchpatch { + name = "remove-deprecation-warning"; + url = "https://github.com/google-deepmind/rlax/commit/dea6eb479ffc32156aefe73015387a762c6b4562.patch"; + hash = "sha256-htDyDRJW0eQx7AmrS3Fl7Lbh2VAmoYiDgHSePsQUaWs="; + }) + (fetchpatch { + name = "fix-deprecation-warnings"; + url = "https://github.com/google-deepmind/rlax/commit/605e0ef8ad8f9a06e88d4aabbb7d50e086d0cf3a.patch"; + hash = "sha256-GZ/nGMXne6Lv6yDm/29NVTWxLBVSzaPYKAfQOLHY4UI="; + }) + # https://github.com/google-deepmind/rlax/pull/135 + (fetchpatch { + name = "fix-jax-0.6.0-compat"; + url = "https://github.com/google-deepmind/rlax/commit/461b4cf9b4239d6b1b83aad6e5946f68d8402b93.patch"; + hash = "sha256-uPMpm4IcoBWJwnyuIRjQEfo0F9HIW/lrwecxGW/Yw38="; + }) ]; build-system = [ diff --git a/pkgs/development/python-modules/waymax/default.nix b/pkgs/development/python-modules/waymax/default.nix index 1e499b92f03a..6fb36ddcf8b8 100644 --- a/pkgs/development/python-modules/waymax/default.nix +++ b/pkgs/development/python-modules/waymax/default.nix @@ -19,9 +19,9 @@ tqdm, }: -buildPythonPackage rec { +buildPythonPackage { pname = "waymax"; - version = "0-unstable-2025-03-25"; + version = "0-unstable-2024-03-23"; pyproject = true; src = fetchFromGitHub { @@ -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 = [