From fcc2520c0bde2c6d37855d439ba1e18abfc0e941 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Fri, 18 Sep 2026 07:11:42 +0000 Subject: [PATCH] python3Packages.jax: 0.11.1 -> 0.11.2 Diff: https://github.com/google/jax/compare/jax-v0.11.1...jax-v0.11.2 Changelog: https://docs.jax.dev/en/latest/changelog.html --- .../jax-cuda12-pjrt/default.nix | 4 ++-- .../jax-cuda12-plugin/default.nix | 12 +++++----- .../jax-cuda13-pjrt/default.nix | 4 ++-- .../jax-cuda13-plugin/default.nix | 12 +++++----- .../python-modules/jax/default.nix | 22 ++----------------- .../development/python-modules/jaxlib/bin.nix | 20 ++++++++--------- 6 files changed, 28 insertions(+), 46 deletions(-) diff --git a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix index 0723a76da0a4..dd45c995d18d 100644 --- a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix @@ -17,11 +17,11 @@ let platforms = { x86_64-linux = { name = "manylinux_2_27_x86_64"; - hash = "sha256-/dnXJAjhf3ojvfombpyQE2YgpyTfWEYp81BnIuKy9DQ="; + hash = "sha256-ITOzr8YzB8AWrkmA8Ways94Zl1k2Pq1099VQ6q6VNGM="; }; aarch64-linux = { name = "manylinux_2_27_aarch64"; - hash = "sha256-QAIdRVvZRSJwTjIdMOitomp8Cm75XEp1QqvaATiinRw="; + hash = "sha256-82JH1AlzsgrPPGnax8nHINqqQJaoCbPo/XApK+YXzes="; }; }; currentPlatform = platforms.${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 51adc910d9ce..5bc26ccee746 100644 --- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix @@ -19,17 +19,17 @@ let x86_64-linux = { name = "manylinux_2_27_x86_64"; hashes = { - cp312 = "sha256-8BVPvyNrsiMab8pSBN0eQ6wMHPDeOgAg//kQPlP7hXI="; - cp313 = "sha256-2O/FH9SZiimQxGfmPDhs/ckYfH0QWCXFxKt0Pn1SLRE="; - cp314 = "sha256-a7jxBoIVF+NEiKUyDkX13vSkYo8qJUr6bix+7DVBm74="; + cp312 = "sha256-rmekzEmzTFdXE6gxH3WX9nAOxaK2r728PbVVD7Aihzw="; + cp313 = "sha256-Ve1tvfTW9XxsewqtrrjjNzcaJSSaDjwe+e/LYHsLvGc="; + cp314 = "sha256-ltGYFX526tkZWHQkZvXHRiYY3SeDHcgjhIgOLoDr5OA="; }; }; aarch64-linux = { name = "manylinux_2_27_aarch64"; hashes = { - cp312 = "sha256-8qOC5dv2YzB0C/TbOQPGEuhLJN4JkTx5U+0kA8KWtoY="; - cp313 = "sha256-6KjWb0PQCzXAQwNmsRfBthKDdk8NENOWWr/jn4coSG8="; - cp314 = "sha256-PKEZswHZqTvLg3XGk3z3lJL/WVHKXLuLRvfjSvQtz9M="; + cp312 = "sha256-rQ+7YFSFhdXw3y/hHrjm1BqYVidO2bw+vi0eHhw9AS8="; + cp313 = "sha256-KRqt6kpSxTo1+VfKWjwN1TFrXAgvy7UfFkt/INTh1lk="; + cp314 = "sha256-7H4G16RQgHln81CR1K9sCnJGG9x3xLIXGKUNg3fq6vs="; }; }; }; diff --git a/pkgs/development/python-modules/jax-cuda13-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda13-pjrt/default.nix index 295e06bd2acb..e5f9becb63e9 100644 --- a/pkgs/development/python-modules/jax-cuda13-pjrt/default.nix +++ b/pkgs/development/python-modules/jax-cuda13-pjrt/default.nix @@ -17,11 +17,11 @@ let platforms = { x86_64-linux = { name = "manylinux_2_27_x86_64"; - hash = "sha256-TKiiYaVIWtuFL96VTkuRRfVCwNivRwSojgX67jVBryE="; + hash = "sha256-5f9YnUMGSRO7dAJ2gtiWOcfqSPn4HCuub8dRF+kT+T4="; }; aarch64-linux = { name = "manylinux_2_27_aarch64"; - hash = "sha256-gZx0TjXd8Cckv0WWpijSvjhzaah82mBXUTfuQ37NeXo="; + hash = "sha256-Y7VeUOXCWDol4Za6BrA0tj0rLYrUwvF5NWUsjTrS3S0="; }; }; currentPlatform = platforms.${stdenv.hostPlatform.system}; diff --git a/pkgs/development/python-modules/jax-cuda13-plugin/default.nix b/pkgs/development/python-modules/jax-cuda13-plugin/default.nix index 1cb6276bf2a2..995dbcb99200 100644 --- a/pkgs/development/python-modules/jax-cuda13-plugin/default.nix +++ b/pkgs/development/python-modules/jax-cuda13-plugin/default.nix @@ -19,17 +19,17 @@ let x86_64-linux = { name = "manylinux_2_27_x86_64"; hashes = { - cp312 = "sha256-2zSLkZmBC+W1OAwRObTUgC1t7G8Mpr8Z/LgYgFNB9ZA="; - cp313 = "sha256-C1ebRy0AbJPOUGORCbdkhfLReJ3k4b+fAOz2eXN9Jf0="; - cp314 = "sha256-6CAG8dgVLjTPUySYWk1tcKHfrSXSWzP6Qhfhh8Z9vNI="; + cp312 = "sha256-rTU+0YdyYQVg36KNXhVNntN3/SVyggpI9bVZEDe4eQk="; + cp313 = "sha256-l5aBOt4ZkL18neR1ubqmqidclzYSOQEb/DICKXZvLV4="; + cp314 = "sha256-hMDOpE60Ju825kzR3RRhmauiFWFyMSOQlpBHwmvM7qM="; }; }; aarch64-linux = { name = "manylinux_2_27_aarch64"; hashes = { - cp312 = "sha256-9udjg5/ChxzqewgrphHNOlufxQhstED7tMkgmPmxtWc="; - cp313 = "sha256-TqmizXyPmKVGVCF1mZW1HntI20yopYFFvOUyYnYC2CA="; - cp314 = "sha256-74XFhZxLy91yNLtvh152zfh0ElzV4eZBfJVzfUzeRNc="; + cp312 = "sha256-AWY4p5REkWPKnN1ZvMvQ44AovyjnMhfLlryLFCpbjOI="; + cp313 = "sha256-TIhgusB4zeTetWV6o4l7A5lM/GYngGiUc1QPL+89Nzw="; + cp314 = "sha256-P0ZWPnmhGL359m+VmMeFrjxtIX74NPE4Qx6OdCZDlPM="; }; }; }; diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index 86b7641d2a37..83c3206c995f 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -44,7 +44,7 @@ let in buildPythonPackage (finalAttrs: { pname = "jax"; - version = "0.11.1"; + version = "0.11.2"; pyproject = true; __structuredAttrs = true; @@ -53,7 +53,7 @@ buildPythonPackage (finalAttrs: { repo = "jax"; # google/jax contains tags for jax and jaxlib. Only use jax tags! tag = "jax-v${finalAttrs.version}"; - hash = "sha256-OiH4qhVK7T6o+lYtP1e2UqtSitxVdzUWC5YXbaNMZsQ="; + hash = "sha256-lBfsbEa5Z6GmGVSgdDKoMF0R1TuU+e05QYP96qPNnxQ="; }; build-system = [ setuptools ]; @@ -129,13 +129,6 @@ buildPythonPackage (finalAttrs: { disabledTests = [ # Exceeds tolerance when the machine is busy "test_custom_linear_solve_aux" - - # pytest-xdist/execnet cannot serialize the numpy `type` objects this test passes to - # self.subTest(dtype=...) when shipping subtest reports between workers. - # The assertions themselves pass; the failure is a harness artifact of running with - # --numprocesses. - # New test in jax 0.10.2 (tests/random_impl_test.py). - "test_random_bits" ] ++ lib.optionals usingMKL [ # See @@ -152,17 +145,6 @@ buildPythonPackage (finalAttrs: { # AssertionError: 'INFO' not found in 'I0905 10:11:44.349869 24500 pjrt_api.cc:119] GetPjrtApi was found for cuda at... "test_subprocess_stderr_info_logging" - ] - ++ lib.optionals stdenv.hostPlatform.isx86_64 [ - # The Mosaic GPU interpreter emulates tcgen05 MMA on the CPU backend and compares the - # result with `assert_array_equal`. On x86_64 the two sides contract differently and - # disagree by a single float32 ULP (max relative difference 5.5e-07). - # Passes on aarch64-linux and aarch64-darwin. - "test_async_copy_tmem_with_mma" - "test_can_commit_mma_to_multiple_barriers" - "test_can_deallocate_tmem_while_mma_active_on_different_tmem" - "test_can_pipeline_with_multiple_children" - "test_can_pipeline_with_multiple_parents" ]; pythonImportsCheck = [ "jax" ]; diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix index 73e3276a2b8c..596e8476705a 100644 --- a/pkgs/development/python-modules/jaxlib/bin.nix +++ b/pkgs/development/python-modules/jaxlib/bin.nix @@ -17,7 +17,7 @@ }: let - version = "0.11.1"; + version = "0.11.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 @@ -47,49 +47,49 @@ let "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp312"; - hash = "sha256-w9zsO7M+v1x23DP5YseWZSwHftskZGT6Mssdj8tgh+Y="; + hash = "sha256-Yzih8JVvYNAlKUBKwzqAcAFaUmNc22u7FdDRmXVwd7s="; }; "3.12-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp312"; - hash = "sha256-R7qzp9qKgY+efd7NTtgVQpyqKEDtqelou20zudjFIGc="; + hash = "sha256-1JcAToDaqRuYJY1kWzdOjf9IC/AX0sXz4wsTcayPVLE="; }; "3.12-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp312"; - hash = "sha256-OE9jMzHAEkkczWXDjUfNgNQTH1s4UmlHZr/6k4XdyfM="; + hash = "sha256-7KgOdFt2IzaadrX2hWduwjOjOlXWn2D6flfL78/Uetw="; }; "3.13-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp313"; - hash = "sha256-fr6Gp7iR/NqD5/Y0pnfShWgae4Cz7J7UNVNgeKg3Gs8="; + hash = "sha256-XzysjXAwwfgKGCAl1xceo7EkUWKnEPbljSXuP+F0iqw="; }; "3.13-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp313"; - hash = "sha256-OmeaSbjTKbbkMfOutmZGfG76x1Djsu6Q0OdO/nlZ5kM="; + hash = "sha256-X504M8sr6l40a9ITi6dJkgF+CEHGPYkKAVehP5y/Q54="; }; "3.13-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp313"; - hash = "sha256-x1GK8xaKtM8KALkvVGMBi68NrT+qOzWKPWiJlq+m/VA="; + hash = "sha256-dmvZDieg/1O4e7nXDHNWXsgLmxFlMgzLj3LeRrhVEN8="; }; "3.14-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp314"; - hash = "sha256-c56teS3bhUyzyl3s/gCqsIpzETGrLL7tm3T0VFyS/QE="; + hash = "sha256-HzHYprH7Exgf8YB3V1BtEPbny6aycH9/gPc7hTyy/yI="; }; "3.14-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp314"; - hash = "sha256-ImYW3wJ/E0jad3ZHU8bB4vTUW1VRtnvP3K/PnK3y4oA="; + hash = "sha256-BhRo61rGtiEyFf0RArI3vLyT8kSAyIEkW3NhnQwqMh4="; }; "3.14-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp314"; - hash = "sha256-oGPWEQ0papNZM8Nap+1hs6Iq4kZuOINO7NSK4tFPuKM="; + hash = "sha256-7a1rKhjWPcKWTMl6FtEm2tQWIe2sedZR4PbIeF2JE2k="; }; }; in