1 files changed, 38 insertions, 25 deletions
diff --git a/pkgs/development/python-modules/blackjax/default.nix b/pkgs/development/python-modules/blackjax/default.nix
index 7ecc4eda95198..3e8c680f7c91c 100644
--- a/pkgs/development/python-modules/blackjax/default.nix
+++ b/pkgs/development/python-modules/blackjax/default.nix
@@ -1,21 +1,23 @@
-{ lib
-, buildPythonPackage
-, pythonOlder
-, fetchFromGitHub
-, pytest-xdist
-, pytestCheckHook
-, setuptools-scm
-, fastprogress
-, jax
-, jaxlib
-, jaxopt
-, optax
-, typing-extensions
+{
+  lib,
+  stdenv,
+  buildPythonPackage,
+  pythonOlder,
+  fetchFromGitHub,
+  pytest-xdist,
+  pytestCheckHook,
+  setuptools-scm,
+  fastprogress,
+  jax,
+  jaxlib,
+  jaxopt,
+  optax,
+  typing-extensions,
 }:
 
 buildPythonPackage rec {
   pname = "blackjax";
-  version = "1.1.1";
+  version = "1.2.1";
   pyproject = true;
 
   disabled = pythonOlder "3.9";
@@ -24,12 +26,12 @@ buildPythonPackage rec {
     owner = "blackjax-devs";
     repo = "blackjax";
     rev = "refs/tags/${version}";
-    hash = "sha256-6+ElY1F8oRCtWT4a/LIG6hYMthlq5mDx2baKAc6zIns=";
+    hash = "sha256-VoWBCjFMyE5LVJyf7du/pKlnvDHj22lguiP6ZUzH9ak=";
   };
 
-  nativeBuildInputs = [ setuptools-scm ];
+  build-system = [ setuptools-scm ];
 
-  propagatedBuildInputs = [
+  dependencies = [
     fastprogress
     jax
     jaxlib
@@ -42,15 +44,26 @@ buildPythonPackage rec {
     pytestCheckHook
     pytest-xdist
   ];
-  disabledTestPaths = [ "tests/test_benchmarks.py" ];
-  disabledTests = [
-    # too slow
-    "test_adaptive_tempered_smc"
-  ];
 
-  pythonImportsCheck = [
-    "blackjax"
-  ];
+  disabledTestPaths =
+    [ "tests/test_benchmarks.py" ]
+    ++ lib.optionals (stdenv.isLinux && stdenv.isAarch64) [
+      # Assertion errors on numerical values
+      "tests/mcmc/test_integrators.py"
+    ];
+
+  disabledTests =
+    [
+      # too slow
+      "test_adaptive_tempered_smc"
+    ]
+    ++ lib.optionals (stdenv.isLinux && stdenv.isAarch64) [
+      # Numerical test (AssertionError)
+      # https://github.com/blackjax-devs/blackjax/issues/668
+      "test_chees_adaptation"
+    ];
+
+  pythonImportsCheck = [ "blackjax" ];
 
   meta = with lib; {
     homepage = "https://blackjax-devs.github.io/blackjax";