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

python312Packages.jax[lib]: 0.5.0 -> 0.5.1 (#384879)

Browse Source
This commit is contained in:
Nick Cao
2025-03-01 16:52:13 -05:00
committed by GitHub
parent 2f0c9948ef 42b3111987
commit 2b0c931ffc
10 changed files with 69 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkgs/development/python-modules/distrax/default.nix
+3
View File
@@ -42,6 +42,9 @@ buildPythonPackage rec {
# Flaky: AssertionError: 1 not less than 0.7000000000000001
"test_von_mises_sample_uniform_ks_test"
# Flaky: AssertionError: Not equal to tolerance
"test_composite_methods_are_consistent__with_jit"
# NotImplementedError: Primitive 'square' does not have a registered inverse.
"test_against_tfp_bijectors_square"
"test_log_dets_square__with_device"
pkgs/development/python-modules/equinox/default.nix
+8
View File
@@ -32,6 +32,14 @@ buildPythonPackage rec {
hash = "sha256-hor2qw+aTL7yhV53E/y5DUwyDEYJA8RPRS39xxa8xcw=";
};
# Relax speed constraints on tests that can fail on busy builders
postPatch = ''
substituteInPlace tests/test_while_loop.py \
--replace-fail "speed < 0.1" "speed < 0.5" \
--replace-fail "speed < 0.5" "speed < 1" \
--replace-fail "speed < 1" "speed < 4" \
'';
build-system = [ hatchling ];
dependencies = [
pkgs/development/python-modules/flax/default.nix
+6
View File
@@ -84,6 +84,12 @@ buildPythonPackage rec {
tensorflow
];
pytestFlagsArray = [
# DeprecationWarning: linear_util.wrap_init is missing a DebugInfo object.
"-W"
"ignore::DeprecationWarning"
];
disabledTestPaths = [
# Docs test, needs extra deps + we're not interested in it.
"docs/_ext/codediff_test.py"
pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
+1 -1
View File
@@ -39,7 +39,7 @@ let
srcs = {
"x86_64-linux" = fetchurl {
url = "https://storage.googleapis.com/jax-releases/cuda12_plugin/jax_cuda12_pjrt-${version}-py3-none-manylinux2014_x86_64.whl";
hash = "sha256-0jgzwbiF2WwnZAAOlQUvK1gnx31JLqaPZ+kDoTJlbbs=";
hash = "sha256-05Xe87NP1oSOEVlu8pdaiV0fUG31EuQbH8XS3lIMjlE=";
};
# "aarch64-linux" = fetchurl {
# url = "https://storage.googleapis.com/jax-releases/cuda12_plugin/jax_cuda12_pjrt-${version}-py3-none-manylinux2014_aarch64.whl";
pkgs/development/python-modules/jax-cuda12-plugin/default.nix
+8 -8
View File
@@ -40,42 +40,42 @@ let
"3.10-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp310";
hash = "sha256-D0Q6azcpjt+weW/NvR+GzoWksIS2vT8fUKT7/Wfe2Gs=";
hash = "sha256-ymCGSgWlzqRK51dthHtHeeTeYYUKmhgjg0H8Q6dY1Vs=";
};
"3.10-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp310";
hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
hash = "sha256-488emSaNinBBPw+sM1krh2nBPckdch+RxYeEa+nYhAM=";
};
"3.11-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp311";
hash = "sha256-qYE1oCIwZLj1xoU+It3BpOOGIVLTf7aF8Nve/+DIASI=";
hash = "sha256-NGVbjq/H2b/sMbB3rBPULgjY7YZV0kFHxa38AVFSaU8=";
};
"3.11-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp311";
hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
hash = "sha256-ik8Mje8QnXb1uqwbKuRpCURFzS1/vYxZf83WI+sC/1Q=";
};
"3.12-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp312";
hash = "sha256-QwWN/FZdjJ2mn0fNTkuVxJXxaG8onvRYTCtygD5vFgc=";
hash = "sha256-Fp22Rbr+whOO2YOvjxTk0RqElyivpXIC55qRBNmJLxY=";
};
"3.12-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp312";
hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
hash = "sha256-wqhtVj3AR5rLTHmuz/vuwYn6vY6XAP4/cxvwOV4dbBE=";
};
"3.13-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp313";
hash = "sha256-3zbEsXbi01qCqfOM13zDadJx5gBR43GgqO9FFD+PWLY=";
hash = "sha256-J0kTjLGHMZBoa7FPMxBIskwXjwXkdIo9L/fSA2c1rT0=";
};
"3.13-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp313";
hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
hash = "sha256-C7aXl3+NbGOS9WLDQF6D68xvXLfCCBQFCE9tXPoJ4yo=";
};
};
in
pkgs/development/python-modules/jax/default.nix
+2 -34
View File
@@ -40,7 +40,7 @@ let
in
buildPythonPackage rec {
pname = "jax";
version = "0.5.0";
version = "0.5.1";
pyproject = true;
src = fetchFromGitHub {
@@ -48,7 +48,7 @@ buildPythonPackage rec {
repo = "jax";
# google/jax contains tags for jax and jaxlib. Only use jax tags!
tag = "jax-v${version}";
hash = "sha256-D6n9Z34nrCbBd9IS8YW6uio5Yi9GLCo9PViO3YYbkQ8=";
hash = "sha256-WXtoLiRGcD8PqCMV+VYGeFr/qFEavuoVM5JSapO1QXc=";
};
build-system = [ setuptools ];
@@ -113,15 +113,6 @@ buildPythonPackage rec {
[
# Exceeds tolerance when the machine is busy
"test_custom_linear_solve_aux"
# UserWarning: Explicitly requested dtype <class 'numpy.float64'>
# requested in astype is not available, and will be truncated to
# dtype float32. (With numpy 1.24)
"testKde3"
"testKde5"
"testKde6"
# Invokes python manually in a subprocess, which does not have the correct dependencies
# ImportError: This version of jax requires jaxlib version >= 0.4.19.
"test_no_log_spam"
]
++ lib.optionals usingMKL [
# See
@@ -133,19 +124,6 @@ buildPythonPackage rec {
"testEigvalsGrad_shape"
]
++ lib.optionals stdenv.hostPlatform.isAarch64 [
# See https://github.com/google/jax/issues/14793.
"test_for_loop_fixpoint_correctly_identifies_loop_varying_residuals_unrolled_for_loop"
"testQdwhWithRandomMatrix3"
"testScanGrad_jit_scan"
# See https://github.com/google/jax/issues/17867.
"test_array"
"test_async"
"test_copy0"
"test_device_put"
"test_make_array_from_callback"
"test_make_array_from_single_device_arrays"
# Fails on some hardware due to some numerical error
# See https://github.com/google/jax/issues/18535
"testQdwhWithOnRankDeficientInput5"
@@ -163,16 +141,6 @@ buildPythonPackage rec {
"test_vmap_mismatched_axis_sizes_error_message_issue_705"
];
disabledTestPaths =
[
# Segmentation fault. See https://gist.github.com/zimbatm/e9b61891f3bcf5e4aaefd13f94344fba
"tests/linalg_test.py"
]
++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
# RuntimeWarning: invalid value encountered in cast
"tests/lax_test.py"
];
pythonImportsCheck = [ "jax" ];
# Test CUDA-enabled jax and jaxlib. Running CUDA-enabled tests is not
pkgs/development/python-modules/jaxlib/bin.nix
+13 -13
View File
@@ -18,7 +18,7 @@
}:
let
version = "0.5.0";
version = "0.5.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
@@ -49,65 +49,65 @@ let
"3.10-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp310";
hash = "sha256-dEQLYyEHM2QA1Pl6Fkgddn8T6pFMU7oU5UTG/aVIGbM=";
hash = "sha256-ZbxJAKBJHfxvubamLoEA0SFCnVjHQplF7CtCTBqCvyc=";
};
"3.10-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp310";
hash = "sha256-Wy7+Pf6/GKhMRR04A6yITuJCAhwRE7J5wT9LvDeMPcA=";
hash = "sha256-CQ/n1LyOGXaHcdvIMZ3pAkdGu1Z/XZujebUQLxdljEE=";
};
"3.10-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp310";
hash = "sha256-G4psQ0XxN/OHZQ3i28SIwgJRt0ErVd1kjhpPE7z1B/s=";
hash = "sha256-LavMsIZHaBj3N80OYi/YjcT1u08N9Melxow0s2bohH8=";
};
"3.11-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp311";
hash = "sha256-CRE+8Vgro018vEQP7bMY9IVbWbd2cRqKuiRzyXJ9MCU=";
hash = "sha256-gMDtVEZkSzg8qj5hdUCAO7DvNuVgnNd1bU0pE6jeUS4=";
};
"3.11-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp311";
hash = "sha256-YwiNv6qFu1bNUhqSWjRy/XMosY7JPC2P+oWvMxCVyZU=";
hash = "sha256-afS54HrQdNRBuZIbeoOu5PT/09VCAz/exQvhRW0GEcY=";
};
"3.11-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp311";
hash = "sha256-bNdi7RYjEySZ+nAcQgNEYQLgqcgsojGUuHKI90bRKik=";
hash = "sha256-M0xJrUEfOaUFXCPxOVUq4yvZr+aWq7scrfnETr72B/c=";
};
"3.12-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp312";
hash = "sha256-+YDHM+mMmYqNqHyajMYbZybQvmZ6WL1mTB1xe0tOrnU=";
hash = "sha256-Uee1n8QLsnBEDFBJs8gvn3/a2uMZnxaBhiDP24C5Z68=";
};
"3.12-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp312";
hash = "sha256-S0sBr7Dd7JbAg1a/8rtoXdvpf9/+Ttbi2DSzCrqXLyI=";
hash = "sha256-W0ulqj9ZtfLjfVJc7dav4P7suIQW5fQ+uacJz97YslA=";
};
"3.12-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp312";
hash = "sha256-c+M1cVdgxW5jUQnWFCZDWl1/RvM2OhFdrqCUJ9XNDv0=";
hash = "sha256-rj3ii/m4Z4HDCjLIi3y9HTIiqMIpqpbL4hBV3NCeuIk=";
};
"3.13-x86_64-linux" = getSrcFromPypi {
platform = "manylinux2014_x86_64";
dist = "cp313";
hash = "sha256-Ee7wHTfA8cUwYmW3byB/EALRNIDe0uMf1j7HaRLJPKI=";
hash = "sha256-3BCf+mhzZAImw2DaeTqKe+r4xIdRrD/bKfprM3kB4YY=";
};
"3.13-aarch64-linux" = getSrcFromPypi {
platform = "manylinux2014_aarch64";
dist = "cp313";
hash = "sha256-fZsXp+oZNV1F7Nsv8NtdcHqG8MWoYtlLibRWjWxFMRo=";
hash = "sha256-3zcE8TXP+H/Z1BkwJIkl8vFjvu1u/qqrzNl0AVgNz9g=";
};
"3.13-aarch64-darwin" = getSrcFromPypi {
platform = "macosx_11_0_arm64";
dist = "cp313";
hash = "sha256-7RjqcWHQOqj9TRtVSUiC8hQg79/qaOXymMSuvPKsPzQ=";
hash = "sha256-jFf7vnmqPOOsLsZXp/F4Z7mzwsrYhbLIOQVnzJc47ug=";
};
};
in
pkgs/development/python-modules/jaxlib/prefetch.sh
+1 -2
View File
@@ -1,7 +1,7 @@
#! /usr/bin/env nix-shell
#! nix-shell -i sh -p jq
prefetch () {
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"
@@ -14,7 +14,6 @@ for py in "310" "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-darwin" "jaxlib-bin"
prefetch "$py" "x86_64-linux" "jax-cuda12-plugin"
prefetch "$py" "aarch64-linux" "jax-cuda12-plugin"
done
pkgs/development/python-modules/numpyro/default.nix
+24 -18
View File
@@ -1,5 +1,6 @@
{
lib,
stdenv,
buildPythonPackage,
fetchFromGitHub,
@@ -75,28 +76,33 @@ buildPythonPackage rec {
"ignore::UserWarning"
];
disabledTests = [
# AssertionError due to tolerance issues
"test_bijective_transforms"
"test_cpu"
"test_entropy_categorical"
"test_gaussian_model"
disabledTests =
[
# AssertionError due to tolerance issues
"test_bijective_transforms"
"test_cpu"
"test_entropy_categorical"
"test_gaussian_model"
# > with pytest.warns(UserWarning, match="Hessian of log posterior"):
# E Failed: DID NOT WARN. No warnings of type (<class 'UserWarning'>,) were emitted.
# E Emitted warnings: [].
"test_laplace_approximation_warning"
# > with pytest.warns(UserWarning, match="Hessian of log posterior"):
# E Failed: DID NOT WARN. No warnings of type (<class 'UserWarning'>,) were emitted.
# E Emitted warnings: [].
"test_laplace_approximation_warning"
# Tests want to download data
"data_load"
"test_jsb_chorales"
# Tests want to download data
"data_load"
"test_jsb_chorales"
# ValueError: compiling computation that requires 2 logical devices, but only 1 XLA devices are available (num_replicas=2)
"test_chain"
# ValueError: compiling computation that requires 2 logical devices, but only 1 XLA devices are available (num_replicas=2)
"test_chain"
# test_biject_to[CorrMatrix()-(15,)] - assert Array(False, dtype=bool)
"test_biject_to"
];
# test_biject_to[CorrMatrix()-(15,)] - assert Array(False, dtype=bool)
"test_biject_to"
]
++ lib.optionals stdenv.hostPlatform.isDarwin [
# AssertionError: Not equal to tolerance rtol=0.06, atol=0
"test_functional_map"
];
meta = {
description = "Library for probabilistic programming with NumPy";
pkgs/development/python-modules/oryx/default.nix
+3
View File
@@ -114,5 +114,8 @@ buildPythonPackage rec {
changelog = "https://github.com/jax-ml/oryx/releases/tag/v${version}";
license = lib.licenses.asl20;
maintainers = with lib.maintainers; [ GaetanLepage ];
# oryx seems to be incompatible with jax 0.5.1
# 237 additional test failures are resulting from the jax bump.
broken = true;
};
}