mirror of
https://github.com/NixOS/nixpkgs.git
synced 2026-07-21 16:11:22 +00:00
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
This commit is contained in:
@@ -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};
|
||||
};
|
||||
|
||||
@@ -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}"
|
||||
|
||||
@@ -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 ];
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user