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

pkgs/development/python-modules/distrax/default.nix

···
     # 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

···
     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

···
     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

···
   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

···
     "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

···
 in
 buildPythonPackage rec {
   pname = "jax";
-  version = "0.5.0";
+  version = "0.5.1";
   pyproject = true;
 
   src = fetchFromGitHub {
···
     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 ];
···
     [
       # 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
···
       "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"
···
       "test_vjp_rule_inconsistent_pytree_structures_error"
       "test_vmap_in_axes_tree_prefix_error"
       "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" ];

pkgs/development/python-modules/jaxlib/bin.nix

···
 }:
 
 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
···
       "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

···
 #! /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"
···
     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

···
 {
   lib,
+  stdenv,
   buildPythonPackage,
   fetchFromGitHub,
 
···
     "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

···
     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;
   };
 }