From 6933478ba61f551b1ee0d08013b81fbfc2240ce3 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Sun, 21 Dec 2025 09:41:20 +0000 Subject: [PATCH 1/5] python3Packages.jaxlib: format prefetch script --- .../python-modules/jaxlib/prefetch.sh | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/pkgs/development/python-modules/jaxlib/prefetch.sh b/pkgs/development/python-modules/jaxlib/prefetch.sh index 269d7823c78a..d27513a46eb2 100755 --- a/pkgs/development/python-modules/jaxlib/prefetch.sh +++ b/pkgs/development/python-modules/jaxlib/prefetch.sh @@ -2,22 +2,22 @@ #! nix-shell -i sh -p jq prefetch() { - expr="(import { system = \"$2\"; config.cudaSupport = true; }).python$1.pkgs.$3.src.url" - url=$(NIX_PATH=.. nix-instantiate --eval -E "$expr" | jq -r) - echo "$url" - sha256=$(nix-prefetch-url "$url") - nix --extra-experimental-features nix-command hash convert --to sri --hash-algo sha256 "$sha256" - echo + expr="(import { system = \"$2\"; config.cudaSupport = true; }).python$1.pkgs.$3.src.url" + url=$(NIX_PATH=.. nix-instantiate --eval -E "$expr" | jq -r) + echo "$url" + sha256=$(nix-prefetch-url "$url") + nix --extra-experimental-features nix-command hash convert --to sri --hash-algo sha256 "$sha256" + echo } for py in "311" "312" "313"; do - prefetch "$py" "x86_64-linux" "jaxlib-bin" - prefetch "$py" "aarch64-linux" "jaxlib-bin" - prefetch "$py" "aarch64-darwin" "jaxlib-bin" - prefetch "$py" "x86_64-linux" "jax-cuda12-plugin" - prefetch "$py" "aarch64-linux" "jax-cuda12-plugin" + prefetch "$py" "x86_64-linux" "jaxlib-bin" + prefetch "$py" "aarch64-linux" "jaxlib-bin" + prefetch "$py" "aarch64-darwin" "jaxlib-bin" + prefetch "$py" "x86_64-linux" "jax-cuda12-plugin" + prefetch "$py" "aarch64-linux" "jax-cuda12-plugin" done for arch in "x86_64-linux" "aarch64-linux"; do - prefetch "312" "$arch" "jax-cuda12-pjrt" + prefetch "312" "$arch" "jax-cuda12-pjrt" done From dad7900f205ec42d91932e2877c47c7a96a5623e Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Sun, 21 Dec 2025 09:38:06 +0000 Subject: [PATCH 2/5] python3Packages.jax: 0.8.1 -> 0.8.2 Diff: https://github.com/google/jax/compare/jax-v0.8.1...jax-v0.8.2 Changelog: https://docs.jax.dev/en/latest/changelog.html --- .../jax-cuda12-pjrt/default.nix | 4 +-- .../jax-cuda12-plugin/default.nix | 22 ++++++++---- .../python-modules/jax/default.nix | 5 +-- .../development/python-modules/jaxlib/bin.nix | 36 +++++++++++++------ .../python-modules/jaxlib/prefetch.sh | 2 +- 5 files changed, 48 insertions(+), 21 deletions(-) diff --git a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix index 763691f76bc7..bf3fe4463dbd 100644 --- a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix @@ -50,8 +50,8 @@ buildPythonPackage rec { .${stdenv.hostPlatform.system}; hash = { - x86_64-linux = "sha256-RStw7hDLmsXX38pV/7zbibbIvGunCkWvfEkNHc6pjrc="; - aarch64-linux = "sha256-pjHQaJkDNUr9ez0uxZW32gamIwp22gD/lUj1QrIbYlA="; + x86_64-linux = "sha256-47q0HKfEjkFj255+/ScbOqhfD+RfXtBwjWu+2TpZ+Xc="; + aarch64-linux = "sha256-cXobGWpkJAnOGV3fAxwgu+rcyIb1Xkmh0/SSc3Ou7a4="; } .${stdenv.hostPlatform.system}; }; diff --git a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix index 08e63bf39784..529692c72c0c 100644 --- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix @@ -39,32 +39,42 @@ let "3.11-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp311"; - hash = "sha256-R5ykONVVAk2sjdEFg3Hvv49HmBmnCuohPz8wN+zpnXQ="; + hash = "sha256-CwozBM5+SUrNjZxZNJDBEqMs22AQ/hr8WE2eQf2GMWc="; }; "3.11-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp311"; - hash = "sha256-szg73AufYmDZrcTKDR9ovyQRWN/mnXJrJnsGgTgup6c="; + hash = "sha256-cNMyIkhK1cN1uPg1e3wjysuET27Pw5Vn+N1H/eboeFg="; }; "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp312"; - hash = "sha256-tgvwu9okzsb6cRcL1pthM1nwGjdtjgn+NL9n7MmjFk8="; + hash = "sha256-IBZYYbPT5m67LA9jpUfR1e4X6kSsO+cVPHkIycqMiPM="; }; "3.12-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp312"; - hash = "sha256-OFAB9W+FKVnwYa4VrRV8OcxEccjR0lRN/D+AVoSsIhM="; + hash = "sha256-QD1eB3MbXNrDvZ+z9Ei9hIAGLLLAq2HqKtI/zQplR5o="; }; "3.13-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp313"; - hash = "sha256-c0LIgQzJR9548oxyh6MLLiAbD1FXhUPdJVNpK3mkmUI="; + hash = "sha256-gsZ5i+Zr+MdzOGkY5MjlzYEZdT87+zyku8RoGCg3UMY="; }; "3.13-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp313"; - hash = "sha256-2nwPLvHGl/mt5Rpxz60hHiv/JUB6aFXd3eNywBkPxGg="; + hash = "sha256-Y3OH3DQIzSBFYmaFAvnpX3bG7d4KbS5I8FUWLcKuvw0="; + }; + "3.14-x86_64-linux" = getSrcFromPypi { + platform = "manylinux_2_27_x86_64"; + dist = "cp314"; + hash = "sha256-pYmLrB2KtgILVFRkQCVkCfLGa8u7OhCZykc8hIQ63a0="; + }; + "3.14-aarch64-linux" = getSrcFromPypi { + platform = "manylinux_2_27_aarch64"; + dist = "cp314"; + hash = "sha256-WMUUc/xiLgMTgDWYX3QYM1ZNcKS9WiF49htizaoy/5Q="; }; }; in diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index efc317e778b5..97876cc1148f 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -41,7 +41,7 @@ let in buildPythonPackage rec { pname = "jax"; - version = "0.8.1"; + version = "0.8.2"; pyproject = true; src = fetchFromGitHub { @@ -49,7 +49,7 @@ buildPythonPackage rec { repo = "jax"; # google/jax contains tags for jax and jaxlib. Only use jax tags! tag = "jax-v${version}"; - hash = "sha256-e1XqcecIl8TGOZmA1FuWg8uwVUvgSdp8uYltXm9cJpQ="; + hash = "sha256-WKdFEhOxJPLjOXOChZbLRGcw0GFeg/TT/FT6M72C6bo="; }; patches = [ @@ -187,6 +187,7 @@ buildPythonPackage rec { meta = { description = "Source-built JAX frontend: differentiate, compile, and transform Numpy code"; homepage = "https://github.com/google/jax"; + changelog = "https://docs.jax.dev/en/latest/changelog.html"; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ GaetanLepage diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix index b31ab8fc648e..42d80878bdf7 100644 --- a/pkgs/development/python-modules/jaxlib/bin.nix +++ b/pkgs/development/python-modules/jaxlib/bin.nix @@ -18,7 +18,7 @@ }: let - version = "0.8.1"; + version = "0.8.2"; 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,49 +49,65 @@ let "3.11-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp311"; - hash = "sha256-IvSJ+1yL4Np75eSVehCTazdgoWlmj4slxdCcUcPvR/Y="; + hash = "sha256-zPd9qReiCTUkfJkGkd7Py90Gwl7wrJTZFKBKrbIvcUw="; }; "3.11-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp311"; - hash = "sha256-/zK2Mg1ykTHvryKTmCW1LXWVfITDKvKwsb2zPPJ7p18="; + hash = "sha256-u4m+RSsbgI0/iPwBxBWzZKJgvkzHrBIMA4AJ9hUKMtw="; }; "3.11-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp311"; - hash = "sha256-hlrdVhOYg0BfPxXJsN5qZKuPSqVJ3/GWty28hr5szB8="; + hash = "sha256-SQvwywKcc8ZclDESS4bNyVCC28H7dvxUnSTXXaM+VFQ="; }; "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp312"; - hash = "sha256-r0kkGJ/FO2kjdxW1bry/xxu5HKFhhBQ9zvDUMMgXPeY="; + hash = "sha256-K5eJvQj4sMxaXBKuiW/kMtWULjLkFwkbi1qWqab9XPE="; }; "3.12-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp312"; - hash = "sha256-vtHpSujHwWvKRHbY1/WC8NGhAqTmnDqb0gaaDcQidKk="; + hash = "sha256-OxblDFtzDJ3QpJ5V8az6pyKwCxrwUipZFVjcwEZCUvI="; }; "3.12-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp312"; - hash = "sha256-iL3g9TXu6maJ4M1X1At2YNUgaslcfULglWKhCbljpJ8="; + hash = "sha256-Aj3m8/Vtoq9wN5cJllAFhjMf21C1MOy7VLlmbaYzvQA="; }; "3.13-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp313"; - hash = "sha256-0kW9aieccspfeW34TN1k18nIq8S42JrfSs9FiY2rlYs="; + hash = "sha256-G/vPbD3iIXhPpM22dloJ1xy0KYsVYms9BAmz382Khmc="; }; "3.13-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp313"; - hash = "sha256-vWl8FxrOHi6dbtkQp484WzxAlc7ikLAlWqWISPKs3qs="; + hash = "sha256-fDBPOgFpZbnR9SOaigOZpzkl9WBP6RTFymbs9zS/ZCI="; }; "3.13-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp313"; - hash = "sha256-oDSfboF53Il9M665DsZrSoBBMw+7uo0HHcYWfNInFTk="; + hash = "sha256-TQBtuWvgIMgWUhKhIWNy+KysT/T4+wZ3Q9aU7yswGs4="; + }; + + "3.14-x86_64-linux" = getSrcFromPypi { + platform = "manylinux_2_27_x86_64"; + dist = "cp314"; + hash = "sha256-5ql9+wIy7tmiu244KOT2gtusGn/qhAv9pXTK4tv1+vk="; + }; + "3.14-aarch64-linux" = getSrcFromPypi { + platform = "manylinux_2_27_aarch64"; + dist = "cp314"; + hash = "sha256-aBCN/w3nStxGgBa+mhn4Dv5IxmDA1aEiKHCUtEsJKvw="; + }; + "3.14-aarch64-darwin" = getSrcFromPypi { + platform = "macosx_11_0_arm64"; + dist = "cp314"; + hash = "sha256-vv+wBOfutcmvskQ54rLPRaTuPj6K30XjVe3yr2Ks+Lg="; }; }; in diff --git a/pkgs/development/python-modules/jaxlib/prefetch.sh b/pkgs/development/python-modules/jaxlib/prefetch.sh index d27513a46eb2..5d7ad4b2c5c5 100755 --- a/pkgs/development/python-modules/jaxlib/prefetch.sh +++ b/pkgs/development/python-modules/jaxlib/prefetch.sh @@ -10,7 +10,7 @@ prefetch() { echo } -for py in "311" "312" "313"; do +for py in "311" "312" "313" "314"; do prefetch "$py" "x86_64-linux" "jaxlib-bin" prefetch "$py" "aarch64-linux" "jaxlib-bin" prefetch "$py" "aarch64-darwin" "jaxlib-bin" From c4a987ff8230cfda6add51bc25eda7058920e0e4 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Fri, 26 Dec 2025 21:17:22 +0000 Subject: [PATCH 3/5] python3Packages.chex: skip failing tests https://github.com/google-deepmind/chex/issues/422 --- pkgs/development/python-modules/chex/default.nix | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/pkgs/development/python-modules/chex/default.nix b/pkgs/development/python-modules/chex/default.nix index b023b751a4cd..dde855d7415c 100644 --- a/pkgs/development/python-modules/chex/default.nix +++ b/pkgs/development/python-modules/chex/default.nix @@ -56,6 +56,16 @@ buildPythonPackage rec { pytestCheckHook ]; + disabledTests = [ + # Jax 0.8.2 incompatibility (reported at https://github.com/google-deepmind/chex/issues/422) + # AssertionError: AssertionError not raised + "test_assert_tree_is_on_device" + # AssertionError: "\[Chex\]\ [\s\S]*sharded arrays are disallowed" does not match ... + "test_assert_tree_is_on_host" + # AssertionError: [Chex] Assertion assert_tree_is_sharded failed: ... + "test_assert_tree_is_sharded" + ]; + meta = { description = "Library of utilities for helping to write reliable JAX code"; homepage = "https://github.com/deepmind/chex"; From 70e989a146b5f8dccb1c95171e08a314dee9b05e Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Fri, 26 Dec 2025 23:48:26 +0000 Subject: [PATCH 4/5] python3Packages.equinox: fix jax 0.8.2 compatibility https://github.com/patrick-kidger/equinox/pull/1162 --- pkgs/development/python-modules/equinox/default.nix | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pkgs/development/python-modules/equinox/default.nix b/pkgs/development/python-modules/equinox/default.nix index 0ede965dc4e7..29337297d713 100644 --- a/pkgs/development/python-modules/equinox/default.nix +++ b/pkgs/development/python-modules/equinox/default.nix @@ -39,6 +39,12 @@ buildPythonPackage rec { --replace-fail "speed < 0.5" "speed < 1" \ --replace-fail "speed < 1" "speed < 20" \ --replace-fail "speed < 2" "speed < 20" + '' + # Fix jax 0.8.2 compat + # Fix submitted upstream: https://github.com/patrick-kidger/equinox/pull/1162 + + '' + substituteInPlace equinox/_ad.py equinox/internal/_primitive.py \ + --replace-fail "jax.core.get_aval(" "jax.typeof(" ''; build-system = [ hatchling ]; From 6b23f6dd9ef6a2780d730b14c4c1757c4f50a38c Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Sun, 28 Dec 2025 17:04:33 +0000 Subject: [PATCH 5/5] python3Packages.array-api-compat: 1.12 -> 1.13 Diff: https://github.com/data-apis/array-api-compat/compare/1.12...1.13 Changelog: https://github.com/data-apis/array-api-compat/releases/tag/1.13 --- .../array-api-compat/default.nix | 41 +++++++++++-------- 1 file changed, 23 insertions(+), 18 deletions(-) diff --git a/pkgs/development/python-modules/array-api-compat/default.nix b/pkgs/development/python-modules/array-api-compat/default.nix index 778fb21a861a..da6af0b98537 100644 --- a/pkgs/development/python-modules/array-api-compat/default.nix +++ b/pkgs/development/python-modules/array-api-compat/default.nix @@ -1,32 +1,37 @@ { lib, + config, buildPythonPackage, fetchFromGitHub, - pytestCheckHook, + + # build-system setuptools, setuptools-scm, - numpy, - jaxlib, - jax, - torch, - dask, - sparse, + + # tests array-api-strict, - config, - cudaSupport ? config.cudaSupport, + dask, + jax, + jaxlib, + numpy, + pytestCheckHook, + sparse, + torch, cupy, + + cudaSupport ? config.cudaSupport, }: buildPythonPackage rec { pname = "array-api-compat"; - version = "1.12"; + version = "1.13"; pyproject = true; src = fetchFromGitHub { owner = "data-apis"; repo = "array-api-compat"; tag = version; - hash = "sha256-Hb0bFjVMl4CBI3gN3abTO2QUPAOvUaFE0GdPjdops5E="; + hash = "sha256-4ZoYtqkY/xPRyBP1xeMR80zMfwiaKtXX/3XzqiweCtc="; }; build-system = [ @@ -35,14 +40,14 @@ buildPythonPackage rec { ]; nativeCheckInputs = [ - pytestCheckHook - numpy - jaxlib - jax - torch - dask - sparse array-api-strict + dask + jax + jaxlib + numpy + pytestCheckHook + sparse + torch ] ++ lib.optionals cudaSupport [ cupy ];