gador/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

python3Packages.jax: 0.8.1 -> 0.8.2 (#472927)

Browse Source
This commit is contained in:
kirillrdy
2025-12-28 22:45:07 +00:00
committed by GitHub
parent 2ccc6ae6eb 6b23f6dd9e
commit e6c252a369
8 changed files with 99 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkgs/development/python-modules/array-api-compat/default.nix
+23 -18
View File
@@ -1,32 +1,37 @@
{
lib,
config,
buildPythonPackage,
fetchFromGitHub,
pytestCheckHook,
# build-system
setuptools,
setuptools-scm,
numpy,
jaxlib,
jax,
torch,
dask,
sparse,
# tests
array-api-strict,
config,
cudaSupport ? config.cudaSupport,
dask,
jax,
jaxlib,
numpy,
pytestCheckHook,
sparse,
torch,
cupy,
cudaSupport ? config.cudaSupport,
}:
buildPythonPackage rec {
pname = "array-api-compat";
version = "1.12";
version = "1.13";
pyproject = true;
src = fetchFromGitHub {
owner = "data-apis";
repo = "array-api-compat";
tag = version;
hash = "sha256-Hb0bFjVMl4CBI3gN3abTO2QUPAOvUaFE0GdPjdops5E=";
hash = "sha256-4ZoYtqkY/xPRyBP1xeMR80zMfwiaKtXX/3XzqiweCtc=";
};
build-system = [
@@ -35,14 +40,14 @@ buildPythonPackage rec {
];
nativeCheckInputs = [
pytestCheckHook
numpy
jaxlib
jax
torch
dask
sparse
array-api-strict
dask
jax
jaxlib
numpy
pytestCheckHook
sparse
torch
]
++ lib.optionals cudaSupport [ cupy ];
pkgs/development/python-modules/chex/default.nix
+10
View File
@@ -56,6 +56,16 @@ buildPythonPackage rec {
pytestCheckHook
];
disabledTests = [
# Jax 0.8.2 incompatibility (reported at https://github.com/google-deepmind/chex/issues/422)
# AssertionError: AssertionError not raised
"test_assert_tree_is_on_device"
# AssertionError: "\[Chex\]\ [\s\S]*sharded arrays are disallowed" does not match ...
"test_assert_tree_is_on_host"
# AssertionError: [Chex] Assertion assert_tree_is_sharded failed: ...
"test_assert_tree_is_sharded"
];
meta = {
description = "Library of utilities for helping to write reliable JAX code";
homepage = "https://github.com/deepmind/chex";
pkgs/development/python-modules/equinox/default.nix
+6
View File
@@ -39,6 +39,12 @@ buildPythonPackage rec {
--replace-fail "speed < 0.5" "speed < 1" \
--replace-fail "speed < 1" "speed < 20" \
--replace-fail "speed < 2" "speed < 20"
''
# Fix jax 0.8.2 compat
# Fix submitted upstream: https://github.com/patrick-kidger/equinox/pull/1162
+ ''
substituteInPlace equinox/_ad.py equinox/internal/_primitive.py \
--replace-fail "jax.core.get_aval(" "jax.typeof("
'';
build-system = [ hatchling ];
pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
+2 -2
View File
@@ -50,8 +50,8 @@ buildPythonPackage rec {
.${stdenv.hostPlatform.system};
hash =
{
x86_64-linux = "sha256-RStw7hDLmsXX38pV/7zbibbIvGunCkWvfEkNHc6pjrc=";
aarch64-linux = "sha256-pjHQaJkDNUr9ez0uxZW32gamIwp22gD/lUj1QrIbYlA=";
x86_64-linux = "sha256-47q0HKfEjkFj255+/ScbOqhfD+RfXtBwjWu+2TpZ+Xc=";
aarch64-linux = "sha256-cXobGWpkJAnOGV3fAxwgu+rcyIb1Xkmh0/SSc3Ou7a4=";
}
.${stdenv.hostPlatform.system};
};
pkgs/development/python-modules/jax-cuda12-plugin/default.nix
+16 -6
View File
@@ -39,32 +39,42 @@ let
"3.11-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp311";
hash = "sha256-R5ykONVVAk2sjdEFg3Hvv49HmBmnCuohPz8wN+zpnXQ=";
hash = "sha256-CwozBM5+SUrNjZxZNJDBEqMs22AQ/hr8WE2eQf2GMWc=";
};
"3.11-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp311";
hash = "sha256-szg73AufYmDZrcTKDR9ovyQRWN/mnXJrJnsGgTgup6c=";
hash = "sha256-cNMyIkhK1cN1uPg1e3wjysuET27Pw5Vn+N1H/eboeFg=";
};
"3.12-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp312";
hash = "sha256-tgvwu9okzsb6cRcL1pthM1nwGjdtjgn+NL9n7MmjFk8=";
hash = "sha256-IBZYYbPT5m67LA9jpUfR1e4X6kSsO+cVPHkIycqMiPM=";
};
"3.12-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp312";
hash = "sha256-OFAB9W+FKVnwYa4VrRV8OcxEccjR0lRN/D+AVoSsIhM=";
hash = "sha256-QD1eB3MbXNrDvZ+z9Ei9hIAGLLLAq2HqKtI/zQplR5o=";
};
"3.13-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp313";
hash = "sha256-c0LIgQzJR9548oxyh6MLLiAbD1FXhUPdJVNpK3mkmUI=";
hash = "sha256-gsZ5i+Zr+MdzOGkY5MjlzYEZdT87+zyku8RoGCg3UMY=";
};
"3.13-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp313";
hash = "sha256-2nwPLvHGl/mt5Rpxz60hHiv/JUB6aFXd3eNywBkPxGg=";
hash = "sha256-Y3OH3DQIzSBFYmaFAvnpX3bG7d4KbS5I8FUWLcKuvw0=";
};
"3.14-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp314";
hash = "sha256-pYmLrB2KtgILVFRkQCVkCfLGa8u7OhCZykc8hIQ63a0=";
};
"3.14-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp314";
hash = "sha256-WMUUc/xiLgMTgDWYX3QYM1ZNcKS9WiF49htizaoy/5Q=";
};
};
in
pkgs/development/python-modules/jax/default.nix
+3 -2
View File
@@ -41,7 +41,7 @@ let
in
buildPythonPackage rec {
pname = "jax";
version = "0.8.1";
version = "0.8.2";
pyproject = true;
src = fetchFromGitHub {
@@ -49,7 +49,7 @@ buildPythonPackage rec {
repo = "jax";
# google/jax contains tags for jax and jaxlib. Only use jax tags!
tag = "jax-v${version}";
hash = "sha256-e1XqcecIl8TGOZmA1FuWg8uwVUvgSdp8uYltXm9cJpQ=";
hash = "sha256-WKdFEhOxJPLjOXOChZbLRGcw0GFeg/TT/FT6M72C6bo=";
};
patches = [
@@ -187,6 +187,7 @@ buildPythonPackage rec {
meta = {
description = "Source-built JAX frontend: differentiate, compile, and transform Numpy code";
homepage = "https://github.com/google/jax";
changelog = "https://docs.jax.dev/en/latest/changelog.html";
license = lib.licenses.asl20;
maintainers = with lib.maintainers; [
GaetanLepage
pkgs/development/python-modules/jaxlib/bin.nix
+26 -10
View File
@@ -18,7 +18,7 @@
}:
let
version = "0.8.1";
version = "0.8.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
@@ -49,49 +49,65 @@ let
"3.11-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp311";
hash = "sha256-IvSJ+1yL4Np75eSVehCTazdgoWlmj4slxdCcUcPvR/Y=";
hash = "sha256-zPd9qReiCTUkfJkGkd7Py90Gwl7wrJTZFKBKrbIvcUw=";
};
"3.11-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp311";
hash = "sha256-/zK2Mg1ykTHvryKTmCW1LXWVfITDKvKwsb2zPPJ7p18=";
hash = "sha256-u4m+RSsbgI0/iPwBxBWzZKJgvkzHrBIMA4AJ9hUKMtw=";
};
"3.11-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp311";
hash = "sha256-hlrdVhOYg0BfPxXJsN5qZKuPSqVJ3/GWty28hr5szB8=";
hash = "sha256-SQvwywKcc8ZclDESS4bNyVCC28H7dvxUnSTXXaM+VFQ=";
};
"3.12-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp312";
hash = "sha256-r0kkGJ/FO2kjdxW1bry/xxu5HKFhhBQ9zvDUMMgXPeY=";
hash = "sha256-K5eJvQj4sMxaXBKuiW/kMtWULjLkFwkbi1qWqab9XPE=";
};
"3.12-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp312";
hash = "sha256-vtHpSujHwWvKRHbY1/WC8NGhAqTmnDqb0gaaDcQidKk=";
hash = "sha256-OxblDFtzDJ3QpJ5V8az6pyKwCxrwUipZFVjcwEZCUvI=";
};
"3.12-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp312";
hash = "sha256-iL3g9TXu6maJ4M1X1At2YNUgaslcfULglWKhCbljpJ8=";
hash = "sha256-Aj3m8/Vtoq9wN5cJllAFhjMf21C1MOy7VLlmbaYzvQA=";
};
"3.13-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp313";
hash = "sha256-0kW9aieccspfeW34TN1k18nIq8S42JrfSs9FiY2rlYs=";
hash = "sha256-G/vPbD3iIXhPpM22dloJ1xy0KYsVYms9BAmz382Khmc=";
};
"3.13-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp313";
hash = "sha256-vWl8FxrOHi6dbtkQp484WzxAlc7ikLAlWqWISPKs3qs=";
hash = "sha256-fDBPOgFpZbnR9SOaigOZpzkl9WBP6RTFymbs9zS/ZCI=";
};
"3.13-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp313";
hash = "sha256-oDSfboF53Il9M665DsZrSoBBMw+7uo0HHcYWfNInFTk=";
hash = "sha256-TQBtuWvgIMgWUhKhIWNy+KysT/T4+wZ3Q9aU7yswGs4=";
};
"3.14-x86_64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_x86_64";
dist = "cp314";
hash = "sha256-5ql9+wIy7tmiu244KOT2gtusGn/qhAv9pXTK4tv1+vk=";
};
"3.14-aarch64-linux" = getSrcFromPypi {
platform = "manylinux_2_27_aarch64";
dist = "cp314";
hash = "sha256-aBCN/w3nStxGgBa+mhn4Dv5IxmDA1aEiKHCUtEsJKvw=";
};
"3.14-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp314";
hash = "sha256-vv+wBOfutcmvskQ54rLPRaTuPj6K30XjVe3yr2Ks+Lg=";
};
};
in
pkgs/development/python-modules/jaxlib/prefetch.sh
+13 -13
View File
@@ -2,22 +2,22 @@
#! nix-shell -i sh -p jq
prefetch() {
expr="(import <nixpkgs> { system = \"$2\"; config.cudaSupport = true; }).python$1.pkgs.$3.src.url"
url=$(NIX_PATH=.. nix-instantiate --eval -E "$expr" | jq -r)
echo "$url"
sha256=$(nix-prefetch-url "$url")
nix --extra-experimental-features nix-command hash convert --to sri --hash-algo sha256 "$sha256"
echo
expr="(import <nixpkgs> { system = \"$2\"; config.cudaSupport = true; }).python$1.pkgs.$3.src.url"
url=$(NIX_PATH=.. nix-instantiate --eval -E "$expr" | jq -r)
echo "$url"
sha256=$(nix-prefetch-url "$url")
nix --extra-experimental-features nix-command hash convert --to sri --hash-algo sha256 "$sha256"
echo
}
for py in "311" "312" "313"; do
prefetch "$py" "x86_64-linux" "jaxlib-bin"
prefetch "$py" "aarch64-linux" "jaxlib-bin"
prefetch "$py" "aarch64-darwin" "jaxlib-bin"
prefetch "$py" "x86_64-linux" "jax-cuda12-plugin"
prefetch "$py" "aarch64-linux" "jax-cuda12-plugin"
for py in "311" "312" "313" "314"; do
prefetch "$py" "x86_64-linux" "jaxlib-bin"
prefetch "$py" "aarch64-linux" "jaxlib-bin"
prefetch "$py" "aarch64-darwin" "jaxlib-bin"
prefetch "$py" "x86_64-linux" "jax-cuda12-plugin"
prefetch "$py" "aarch64-linux" "jax-cuda12-plugin"
done
for arch in "x86_64-linux" "aarch64-linux"; do
prefetch "312" "$arch" "jax-cuda12-pjrt"
prefetch "312" "$arch" "jax-cuda12-pjrt"
done