From 987640768a606967e8080cdf491da9d54f61e067 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Mon, 24 Feb 2025 23:06:21 +0100 Subject: [PATCH 1/6] python312Packages.jax[lib]: 0.5.0 -> 0.5.1 Diff: https://github.com/jax-ml/jax/compare/jax-v0.5.0...jax-v0.5.1 Changelog: https://github.com/jax-ml/jax/releases/tag/jax-v0.5.1 --- .../jax-cuda12-pjrt/default.nix | 2 +- .../jax-cuda12-plugin/default.nix | 16 ++++----- .../python-modules/jax/default.nix | 36 ++----------------- .../development/python-modules/jaxlib/bin.nix | 26 +++++++------- .../python-modules/jaxlib/prefetch.sh | 3 +- 5 files changed, 25 insertions(+), 58 deletions(-) diff --git a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix index d0881e760519..6dff1eb69ef7 100644 --- a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix @@ -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"; diff --git a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix index b557c2597332..e9d22d155f69 100644 --- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix @@ -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 diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index c7822e24e88c..824a8259034d 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -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 - # 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 diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix index 7cc8578a01d3..f969f398edae 100644 --- a/pkgs/development/python-modules/jaxlib/bin.nix +++ b/pkgs/development/python-modules/jaxlib/bin.nix @@ -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 diff --git a/pkgs/development/python-modules/jaxlib/prefetch.sh b/pkgs/development/python-modules/jaxlib/prefetch.sh index 22c5253ab0d8..1a4f38daa6a7 100755 --- a/pkgs/development/python-modules/jaxlib/prefetch.sh +++ b/pkgs/development/python-modules/jaxlib/prefetch.sh @@ -1,7 +1,7 @@ #! /usr/bin/env nix-shell #! nix-shell -i sh -p jq -prefetch () { +prefetch() { expr="(import { 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 From 64250c658389655adf1ebe1b2a88916ffd980291 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Tue, 25 Feb 2025 17:34:04 +0100 Subject: [PATCH 2/6] python312Packages.flax: ignore DeprecationWarning during tests --- pkgs/development/python-modules/flax/default.nix | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix index e3e2d4eeb844..35d3c05e5762 100644 --- a/pkgs/development/python-modules/flax/default.nix +++ b/pkgs/development/python-modules/flax/default.nix @@ -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" From d6da15c15f05ccd6b9eb9bc536609ac6e98ca8d6 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Thu, 27 Feb 2025 16:08:03 +0100 Subject: [PATCH 3/6] python312Packages.equinox: relax speed constraints on tests that can fail on busy builders --- pkgs/development/python-modules/equinox/default.nix | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pkgs/development/python-modules/equinox/default.nix b/pkgs/development/python-modules/equinox/default.nix index 9082cdc78aef..d9bbc922431a 100644 --- a/pkgs/development/python-modules/equinox/default.nix +++ b/pkgs/development/python-modules/equinox/default.nix @@ -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 = [ From 7c871e5b6e1a4c056022f527a07b5cff0830e650 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Thu, 27 Feb 2025 17:18:49 +0100 Subject: [PATCH 4/6] python312Packages.numpyro: disable failing test on darwin --- .../python-modules/numpyro/default.nix | 42 +++++++++++-------- 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/pkgs/development/python-modules/numpyro/default.nix b/pkgs/development/python-modules/numpyro/default.nix index 42962f2e180e..07d928fad232 100644 --- a/pkgs/development/python-modules/numpyro/default.nix +++ b/pkgs/development/python-modules/numpyro/default.nix @@ -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 (,) 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 (,) 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"; From a24be4d8d570bd7b4e440abc4cc61ef74f5733c5 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Sat, 1 Mar 2025 15:11:35 +0100 Subject: [PATCH 5/6] python312Packages.distrax: disable flaky test --- pkgs/development/python-modules/distrax/default.nix | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pkgs/development/python-modules/distrax/default.nix b/pkgs/development/python-modules/distrax/default.nix index c01584c3ab4e..6c9d585a09f5 100644 --- a/pkgs/development/python-modules/distrax/default.nix +++ b/pkgs/development/python-modules/distrax/default.nix @@ -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" From 42b311198748f6d4296efbdbd0d90b9f87a99524 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Sat, 1 Mar 2025 22:40:25 +0100 Subject: [PATCH 6/6] python312Packages.oryx: mark as broken --- pkgs/development/python-modules/oryx/default.nix | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pkgs/development/python-modules/oryx/default.nix b/pkgs/development/python-modules/oryx/default.nix index 169902c8a5be..07f0302104df 100644 --- a/pkgs/development/python-modules/oryx/default.nix +++ b/pkgs/development/python-modules/oryx/default.nix @@ -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; }; }