python3Packages.jax: 0.10.2 -> 0.11.0 (#543317)

This commit is contained in:
Nick Cao
2026-07-20 13:47:47 +00:00
committed by GitHub
6 changed files with 25 additions and 49 deletions

View File

@@ -28,7 +28,7 @@
buildPythonPackage (finalAttrs: {
pname = "blackjax";
version = "1.6";
version = "1.6.2";
pyproject = true;
__structuredAttrs = true;
@@ -36,7 +36,7 @@ buildPythonPackage (finalAttrs: {
owner = "blackjax-devs";
repo = "blackjax";
tag = finalAttrs.version;
hash = "sha256-qLOAmUQxr1xtlJB/TMnjFkvvHUwh0XKpPN+FVD8ju8Y=";
hash = "sha256-NO/CvYtxfAid3ETpj5DcNQPdARP2cwqy9p0kHOybvNg=";
};
build-system = [

View File

@@ -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};
};

View File

@@ -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}"

View File

@@ -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 ];

View File

@@ -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

View File

@@ -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"