From 68bc3961e45980e50925f72a700bb904757aa394 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Sat, 18 Jul 2026 21:02:02 +0000 Subject: [PATCH] python3Packages.jax: 0.10.2 -> 0.11.0 Diff: https://github.com/google/jax/compare/jax-v0.10.2...jax-v0.11.0 Changelog: https://docs.jax.dev/en/latest/changelog.html --- .../jax-cuda12-pjrt/default.nix | 5 +-- .../jax-cuda12-plugin/default.nix | 23 ++++-------- .../python-modules/jax/default.nix | 4 +-- .../development/python-modules/jaxlib/bin.nix | 36 ++++++------------- .../python-modules/jaxlib/prefetch.sh | 2 +- 5 files changed, 23 insertions(+), 47 deletions(-) diff --git a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix index 9b640d02fbcb..dd9c03ed8e75 100644 --- a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix @@ -35,6 +35,7 @@ buildPythonPackage (finalAttrs: { pname = "jax-cuda12-pjrt"; inherit version; pyproject = false; + __structuredAttrs = true; src = fetchPypi { pname = "jax_cuda12_pjrt"; @@ -50,8 +51,8 @@ buildPythonPackage (finalAttrs: { .${stdenv.hostPlatform.system}; hash = { - x86_64-linux = "sha256-gG0f0pA4tqz1orKJ3WIZKr6pd+4mrvYOopWx0oojrPg="; - aarch64-linux = "sha256-tvLWZUjuHukQqDa1/j0Ou+HaBKYtD0aLpIIhiQNqMrM="; + x86_64-linux = "sha256-JC35nJCCeik33xxESYPUB2Fz+YNtwLwdIN+GlgFE81o="; + aarch64-linux = "sha256-Z+xZFefklHddXcDXPR4zy3Ske2prsQie73tRfk6HPjM="; } .${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 01eb6ec63239..928382219e47 100644 --- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix @@ -36,45 +36,35 @@ let # upstream does not distribute jax-cuda12-plugin 0.4.38 binaries for aarch64-linux srcs = { - "3.11-x86_64-linux" = getSrcFromPypi { - platform = "manylinux_2_27_x86_64"; - dist = "cp311"; - hash = "sha256-drv+zi9u/7TokI64pfvxq7wRJd5mdoPsL/ulHGPcV68="; - }; - "3.11-aarch64-linux" = getSrcFromPypi { - platform = "manylinux_2_27_aarch64"; - dist = "cp311"; - hash = "sha256-nNCTcHC70o5EuRQ1tSxyJcjtbLN922qQ6cJGhzoD+c0="; - }; "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp312"; - hash = "sha256-Trbo4Jkv2Yl9skaNoz8nauQU7ITcQ2gEEkQEQwUC7qw="; + hash = "sha256-Mf/iRdDi/Rz07lclT/xMiFBbuOf1639rOjfWkdc2ZQo="; }; "3.12-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp312"; - hash = "sha256-dnoUgtv2UmiEA8TiL/AUcOGt0TxI+dgXFp/Y90QK/OU="; + hash = "sha256-PpXLnerCmI4VImxtiSfnK//IQVBdI2Ix5IQ6ixZnnyU="; }; "3.13-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp313"; - hash = "sha256-lIqYiSfLEIQ7UBohUYHMsdrtP6cFMbLxrghC/hwIsF4="; + hash = "sha256-tvcG1sEP6b+v4DNMXVLkDm6Z80PCTY508EWsZ7f3rOI="; }; "3.13-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp313"; - hash = "sha256-aXdqyEkRL0/C1vphb4dyEG2v84/NuCXE459H6HjydhA="; + hash = "sha256-7xA7NTZHlXCMSkOAN22Nc/w9SHNWpwN/zyDKtzX7WaY="; }; "3.14-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp314"; - hash = "sha256-DHoCBBVVhcwcT8swxmt/iBs45X4fDY6X1I4WVL8NJ/g="; + hash = "sha256-S2sGD+c//eIG5iwmNweWGMNMK0t0vlywNxsipDqQfEo="; }; "3.14-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp314"; - hash = "sha256-SUbsV923leLagojOsc7d1YwqIEoPd2WvierkhqHTrOc="; + hash = "sha256-L43Mb/NIMxWCUbUrcRo7XHN4XsNkdxKUvDgnqq6nVFA="; }; }; in @@ -82,6 +72,7 @@ buildPythonPackage { pname = "jax-cuda12-plugin"; inherit version; pyproject = false; + __structuredAttrs = true; src = ( srcs."${python.pythonVersion}-${stdenv.hostPlatform.system}" diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index cc3ae7a61914..72c78f4f67f1 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -40,7 +40,7 @@ let in buildPythonPackage (finalAttrs: { pname = "jax"; - version = "0.10.2"; + version = "0.11.0"; pyproject = true; __structuredAttrs = true; @@ -49,7 +49,7 @@ buildPythonPackage (finalAttrs: { repo = "jax"; # google/jax contains tags for jax and jaxlib. Only use jax tags! tag = "jax-v${finalAttrs.version}"; - hash = "sha256-OQkh9uC8NsxoG3SByPybXQ81c11T3lYgjaU3tbB0+6E="; + hash = "sha256-EE4JuiiwgdQlTsX6dE8KRjcGZHRiQVDXlDVFHchfyYs="; }; build-system = [ setuptools ]; diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix index 7b4324041314..4b5a1742e6c2 100644 --- a/pkgs/development/python-modules/jaxlib/bin.nix +++ b/pkgs/development/python-modules/jaxlib/bin.nix @@ -18,7 +18,7 @@ }: let - version = "0.10.2"; + version = "0.11.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 @@ -46,68 +46,52 @@ let }; in { - "3.11-x86_64-linux" = getSrcFromPypi { - platform = "manylinux_2_27_x86_64"; - dist = "cp311"; - hash = "sha256-H6yjxdRmLLSmEwpoEF1ou1IHZIF+Fl1u6/1nhsDR8w8="; - }; - "3.11-aarch64-linux" = getSrcFromPypi { - platform = "manylinux_2_27_aarch64"; - dist = "cp311"; - hash = "sha256-1EVl3P0bT2D3bZEcZRIRiopPx2S975JmP+y4v8zVTyM="; - }; - "3.11-aarch64-darwin" = getSrcFromPypi { - platform = "macosx_11_0_arm64"; - dist = "cp311"; - hash = "sha256-WpiHP8hnYjuB8r7hXVVLjt1liKGD0B+lDSGx49uW/ys="; - }; - "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp312"; - hash = "sha256-/ojsRDcUxDeZaLbBCfn6YXx60ZuAKCjk17+GHNZtpLc="; + hash = "sha256-fqncBtlPU23qvPMEUTFDvIn1EXPd/XhtRPisMD799kM="; }; "3.12-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp312"; - hash = "sha256-U7cpd65YLAOp6OHN7h77+OvBQYJwllsOaereV6z0AzE="; + hash = "sha256-G3Z8TXrU38Y1gxO2OubYPDVx086WbeBNiq8p60RjN4Y="; }; "3.12-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp312"; - hash = "sha256-R7t8ARUV6oYr5+gxP0D5xWy+wJ3Jig/LUBZ4X81FTAE="; + hash = "sha256-9b8LO/wu9HeFgARzidrLUa4ADlDGs7b5jRB4iScGbJc="; }; "3.13-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp313"; - hash = "sha256-nkgYtKh1b9ORh2bKKqU0ISWAn08Ipv5GAm1DhufCNkQ="; + hash = "sha256-M92UBcqwXuTz7MOpEokyPpvS8I0lmQ5dRf+4Uk9RlQY="; }; "3.13-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp313"; - hash = "sha256-RbKLAjhperdLvPIEEar7bbQqzDGDbML9cR5c8Fa/lVY="; + hash = "sha256-ESoBxYRwen0OhjT9Vd1+//6BKWbmDC06yE6MJORXEzY="; }; "3.13-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp313"; - hash = "sha256-TfUwr6NUoi3BdHpdVgZARQy7iV1JiJM4o/WMdqTHbI4="; + hash = "sha256-iLNJqpXkUsqjCmcjMg6T7HvRqySenwvykADaG1765zM="; }; "3.14-x86_64-linux" = getSrcFromPypi { platform = "manylinux_2_27_x86_64"; dist = "cp314"; - hash = "sha256-yjTzYxl/sKxAglgsp1UAeRA2njP4qLo9Ne2UtxBwEH0="; + hash = "sha256-YSXahTJhBkG30+p5JiF82rUtzLKU72YVRV4rbsueLuY="; }; "3.14-aarch64-linux" = getSrcFromPypi { platform = "manylinux_2_27_aarch64"; dist = "cp314"; - hash = "sha256-8Y9W/ukGmc+6m2YnBFp6KZcCyw4q+CzhgNmmp8gEgJM="; + hash = "sha256-SDQpzH+n/Wogy1D2ZsjS1Jnvt+m6MWOBHAHwUa1gV8Q="; }; "3.14-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp314"; - hash = "sha256-cuuiixL+4CYW+kKqS4gbSrYtd1fHhDxGJAHT+zSie+Q="; + hash = "sha256-fmwOI3SUO6cZH+GW4iBnt4mJTFjPbd0kyZca5kLeWO4="; }; }; in diff --git a/pkgs/development/python-modules/jaxlib/prefetch.sh b/pkgs/development/python-modules/jaxlib/prefetch.sh index 5d7ad4b2c5c5..5676a832c88f 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" "314"; do +for py in "312" "313" "314"; do prefetch "$py" "x86_64-linux" "jaxlib-bin" prefetch "$py" "aarch64-linux" "jaxlib-bin" prefetch "$py" "aarch64-darwin" "jaxlib-bin"