nathan/nixpkgs
1
0
Fork 0
mirror of https://github.com/NixOS/nixpkgs.git synced 2026-10-30 16:10:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

python3Packages.jax: 0.11.2 -> 0.12.0.1 (#576430)

Browse Source
This commit is contained in:
Gaétan Lepage
2026-10-22 23:21:31 +00:00
committed by GitHub
parent 9cfab8a252 affcb7df7d
commit aeb0b10e84
9 changed files with 45 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
View File

@@ -17,11 +17,11 @@ let
platforms = {
x86_64-linux = {
name = "manylinux_2_27_x86_64";
hash = "sha256-ITOzr8YzB8AWrkmA8Ways94Zl1k2Pq1099VQ6q6VNGM=";
hash = "sha256-XsKOBhmNCUAxw3WQtOp1OujN5IeatLYLr0Dw3/tnMoY=";
};
aarch64-linux = {
name = "manylinux_2_27_aarch64";
hash = "sha256-82JH1AlzsgrPPGnax8nHINqqQJaoCbPo/XApK+YXzes=";
hash = "sha256-5qpsWMVCDBvaItsxovbFwdCmbLcB7rcx/q5E/vGXcXw=";
};
};
currentPlatform = platforms.${stdenv.hostPlatform.system};

12
pkgs/development/python-modules/jax-cuda12-plugin/default.nix
View File

@@ -19,17 +19,17 @@ let
x86_64-linux = {
name = "manylinux_2_27_x86_64";
hashes = {
cp312 = "sha256-rmekzEmzTFdXE6gxH3WX9nAOxaK2r728PbVVD7Aihzw=";
cp313 = "sha256-Ve1tvfTW9XxsewqtrrjjNzcaJSSaDjwe+e/LYHsLvGc=";
cp314 = "sha256-ltGYFX526tkZWHQkZvXHRiYY3SeDHcgjhIgOLoDr5OA=";
cp312 = "sha256-YdxnqDQa1qC7WWMTB5FFga0WBV9sVjWit+V9//t35GY=";
cp313 = "sha256-Vn3g7Q5wiq6vmIULqGmxVtkbIBRRlosGvu/DMdpJMp8=";
cp314 = "sha256-29SW/86v1w3lRyZciw/vbjsJU/WbQKYr3qJMZxUm4Qk=";
};
};
aarch64-linux = {
name = "manylinux_2_27_aarch64";
hashes = {
cp312 = "sha256-rQ+7YFSFhdXw3y/hHrjm1BqYVidO2bw+vi0eHhw9AS8=";
cp313 = "sha256-KRqt6kpSxTo1+VfKWjwN1TFrXAgvy7UfFkt/INTh1lk=";
cp314 = "sha256-7H4G16RQgHln81CR1K9sCnJGG9x3xLIXGKUNg3fq6vs=";
cp312 = "sha256-qIHb70FjEPUC6sY3ZU+and5/5LMZ3lehg4dfiG3WKsc=";
cp313 = "sha256-4pwRBYRMHaWajYHApfF9lNKFJIF7QOHJITwEUAbZThc=";
cp314 = "sha256-5q/fhhyOuwAiYWYAXD2VaO4GM0npGnR+1gCyxEBb1Sc=";
};
};
};

4
pkgs/development/python-modules/jax-cuda13-pjrt/default.nix
View File

@@ -17,11 +17,11 @@ let
platforms = {
x86_64-linux = {
name = "manylinux_2_27_x86_64";
hash = "sha256-5f9YnUMGSRO7dAJ2gtiWOcfqSPn4HCuub8dRF+kT+T4=";
hash = "sha256-WJm7YXftqG6lMq+5omeU8r3QSWb2oB+QKCZrFsK2z5A=";
};
aarch64-linux = {
name = "manylinux_2_27_aarch64";
hash = "sha256-Y7VeUOXCWDol4Za6BrA0tj0rLYrUwvF5NWUsjTrS3S0=";
hash = "sha256-QHWNBC34wGmRtqGoAWtW6c77HtTGd8BiElw+txw9q5c=";
};
};
currentPlatform = platforms.${stdenv.hostPlatform.system};

12
pkgs/development/python-modules/jax-cuda13-plugin/default.nix
View File

@@ -19,17 +19,17 @@ let
x86_64-linux = {
name = "manylinux_2_27_x86_64";
hashes = {
cp312 = "sha256-rTU+0YdyYQVg36KNXhVNntN3/SVyggpI9bVZEDe4eQk=";
cp313 = "sha256-l5aBOt4ZkL18neR1ubqmqidclzYSOQEb/DICKXZvLV4=";
cp314 = "sha256-hMDOpE60Ju825kzR3RRhmauiFWFyMSOQlpBHwmvM7qM=";
cp312 = "sha256-YCj68XnnqnzGQQDtPHx7GxE0b1qgXV6eM+j+YMUon6E=";
cp313 = "sha256-SBcsMa7srjEhwFYbVNm4otDHVUk1l5jc7/PluevPdN8=";
cp314 = "sha256-FrfbYs0i21END/Uv268dOxSQic/2oZPRldkinY6xQ3k=";
};
};
aarch64-linux = {
name = "manylinux_2_27_aarch64";
hashes = {
cp312 = "sha256-AWY4p5REkWPKnN1ZvMvQ44AovyjnMhfLlryLFCpbjOI=";
cp313 = "sha256-TIhgusB4zeTetWV6o4l7A5lM/GYngGiUc1QPL+89Nzw=";
cp314 = "sha256-P0ZWPnmhGL359m+VmMeFrjxtIX74NPE4Qx6OdCZDlPM=";
cp312 = "sha256-CH8TMc1asn3NQwO61xKe+vfx3u2DDVZLjQXX7+yUbWE=";
cp313 = "sha256-hGeGn/1kvpd7DzwiuY/B+NYBl3Z8ZI9qLHv7/1ZIlK4=";
cp314 = "sha256-in5K+eDFVwsue9VCOva2hm26pxndlR0kF7lIBvfJFCs=";
};
};
};

4
pkgs/development/python-modules/jax/default.nix
View File

@@ -44,7 +44,7 @@ let
in
buildPythonPackage (finalAttrs: {
pname = "jax";
version = "0.11.2";
version = "0.12.0.1";
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-lBfsbEa5Z6GmGVSgdDKoMF0R1TuU+e05QYP96qPNnxQ=";
hash = "sha256-um7ZO79sZCVFW0Ci4igjabBXildH698V1mCKpIBUVe8=";
};
build-system = [ setuptools ];

20
pkgs/development/python-modules/jaxlib/bin.nix
View File

@@ -17,7 +17,7 @@
}:
let
version = "0.11.2";
version = "0.12.0.1";
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-Yzih8JVvYNAlKUBKwzqAcAFaUmNc22u7FdDRmXVwd7s=";
hash = "sha256-eDGOmseckpcvgcdwkPVXvbSlpJghdRxu4/apxmqUiEI=";
};
"3.12-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp312";
hash = "sha256-1JcAToDaqRuYJY1kWzdOjf9IC/AX0sXz4wsTcayPVLE=";
hash = "sha256-QIEiAcO81MsQS4j3H1j+xMvQwyupDlySs/qPQIGfZ3Q=";
};
"3.12-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp312";
hash = "sha256-7KgOdFt2IzaadrX2hWduwjOjOlXWn2D6flfL78/Uetw=";
hash = "sha256-2yzIWE19DRZik+8kdbcmmIoQ9vm1qOO443EkFeVXA4A=";
};
"3.13-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp313";
hash = "sha256-XzysjXAwwfgKGCAl1xceo7EkUWKnEPbljSXuP+F0iqw=";
hash = "sha256-/cVpmbppC9kV9wKWEyaYSOqvhdIx7Iq8RmENO4QMGJw=";
};
"3.13-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp313";
hash = "sha256-X504M8sr6l40a9ITi6dJkgF+CEHGPYkKAVehP5y/Q54=";
hash = "sha256-yN4crTgnExFPZX9AY4PtxT7p/aRnnbTisJ0d/v4cTkU=";
};
"3.13-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp313";
hash = "sha256-dmvZDieg/1O4e7nXDHNWXsgLmxFlMgzLj3LeRrhVEN8=";
hash = "sha256-i+FgsLhHvFcWjkMFcTdW9toNC7omcV5k+dL0bAp6a0I=";
};
"3.14-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp314";
hash = "sha256-HzHYprH7Exgf8YB3V1BtEPbny6aycH9/gPc7hTyy/yI=";
hash = "sha256-tgeadmnWJOvULRVxM5Vm5iz+ciX0taILjhZ3PqmoNXo=";
};
"3.14-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp314";
hash = "sha256-BhRo61rGtiEyFf0RArI3vLyT8kSAyIEkW3NhnQwqMh4=";
hash = "sha256-YJ3v+uo4jR7Z0cUQTY9gQXkfAHca3DBfcxa3QvUC/Ws=";
};
"3.14-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp314";
hash = "sha256-7a1rKhjWPcKWTMl6FtEm2tQWIe2sedZR4PbIeF2JE2k=";
hash = "sha256-k7jaGD8CNWA8RgtTDajNbrI+p/HGrcT4KkOmA1SX1GY=";
};
};
in

4
pkgs/development/python-modules/jaxopt/default.nix
View File

@@ -25,6 +25,7 @@ buildPythonPackage (finalAttrs: {
pname = "jaxopt";
version = "0.8.5";
pyproject = true;
__structuredAttrs = true;
src = fetchFromGitHub {
owner = "google";
@@ -88,6 +89,9 @@ buildPythonPackage (finalAttrs: {
# https://github.com/google/jaxopt/issues/618
"test_binary_logit_log_likelihood"
# AssertionError: Array(False, dtype=bool) is not true (since jax 0.12.0)
"test_hager_zhang_linesearch_non_finite1"
# AssertionError (flaky numerical tests)
"test_Rosenbrock2"
"test_Rosenbrock5"

1
pkgs/development/python-modules/numpyro/default.nix
View File

@@ -102,6 +102,7 @@ buildPythonPackage (finalAttrs: {
"test_cpu"
"test_entropy_categorical"
"test_gaussian_model"
"test_zero_inflated_logits_probs_agree"
# > with pytest.warns(UserWarning, match="Hessian of log posterior"):
# E Failed: DID NOT WARN. No warnings of type (<class 'UserWarning'>,) were emitted.

12
pkgs/development/python-modules/waymax/default.nix
View File

@@ -38,6 +38,18 @@ buildPythonPackage (finalAttrs: {
hash = "sha256-vS0ZxLPKcjRzu5QaFApOVYNO5LBNQjop8wzmdEVdMC8=";
};
# jax>=0.12 no longer supports 2D vectors in jnp.cross
# ValueError: Dimension must be 3 for cross product
postPatch = ''
substituteInPlace waymax/metrics/roadgraph.py \
--replace-fail \
"jnp.cross(points_to_edge, nearest_vector_xys)" \
"(points_to_edge[..., 0] * nearest_vector_xys[..., 1] - points_to_edge[..., 1] * nearest_vector_xys[..., 0])" \
--replace-fail \
"jnp.cross(points_to_edge, prior_vector_xys)" \
"(points_to_edge[..., 0] * prior_vector_xys[..., 1] - points_to_edge[..., 1] * prior_vector_xys[..., 0])"
'';
build-system = [ setuptools ];
dependencies = [