From 06e0c1e54b302235b6bf4c5916ab346ced4aa93d Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Thu, 26 Mar 2026 20:52:53 +0000
Subject: [PATCH 1/2] python3Packages.blackjax: 1.3 -> 1.4

Diff: https://github.com/blackjax-devs/blackjax/compare/1.3...1.4

Changelog: https://github.com/blackjax-devs/blackjax/releases/tag/1.4
---
 .../python-modules/blackjax/default.nix       | 26 +++++++++----------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/pkgs/development/python-modules/blackjax/default.nix b/pkgs/development/python-modules/blackjax/default.nix
index 5077178f3b11..1e40b9ac7d9e 100644
--- a/pkgs/development/python-modules/blackjax/default.nix
+++ b/pkgs/development/python-modules/blackjax/default.nix
@@ -4,14 +4,17 @@
   fetchFromGitHub,
 
   # build-system
+  setuptools,
   setuptools-scm,
 
   # dependencies
   fastprogress,
+  ipython,
   jax,
   jaxlib,
-  jaxopt,
+  numpy,
   optax,
+  scipy,
   typing-extensions,
 
   # checks
@@ -22,27 +25,29 @@
 
 buildPythonPackage (finalAttrs: {
   pname = "blackjax";
-  version = "1.3";
+  version = "1.4";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "blackjax-devs";
     repo = "blackjax";
     tag = finalAttrs.version;
-    hash = "sha256-ystvPfIsnMFYkC+LNtcRQsI19i/y/905SnPSApM8v4E=";
+    hash = "sha256-x/K/M7C+XNhqMdRZdDPpKmGgmnrjGsruDL3lFia2ioQ=";
   };
 
-  build-system = [ setuptools-scm ];
-
-  pythonRelaxDeps = [
-    "jaxopt"
+  build-system = [
+    setuptools
+    setuptools-scm
   ];
+
   dependencies = [
     fastprogress
+    ipython
     jax
     jaxlib
-    jaxopt
+    numpy
     optax
+    scipy
     typing-extensions
   ];
 
@@ -52,11 +57,6 @@ buildPythonPackage (finalAttrs: {
     pytest-xdist
   ];
 
-  pytestFlags = [
-    # DeprecationWarning: JAXopt is no longer maintained
-    "-Wignore::DeprecationWarning"
-  ];
-
   disabledTestPaths = [
     "tests/test_benchmarks.py"
 

From 49c11d92142613f4fb617988db8cca39ce11a7a9 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Thu, 26 Mar 2026 22:32:19 +0000
Subject: [PATCH 2/2] python3Packages.bambi: skip failing tests

---
 pkgs/development/python-modules/bambi/default.nix | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pkgs/development/python-modules/bambi/default.nix b/pkgs/development/python-modules/bambi/default.nix
index 5f96a8f49cb1..dc271a364f8e 100644
--- a/pkgs/development/python-modules/bambi/default.nix
+++ b/pkgs/development/python-modules/bambi/default.nix
@@ -85,6 +85,12 @@ buildPythonPackage (finalAttrs: {
     # https://github.com/bambinos/bambi/issues/888
     "test_beta_regression"
 
+    # Failing since blackjax was updated to 1.4
+    # ValueError: cannot select an axis to squeeze out which has size not equal to one,
+    # got shape=(4, 2) and dimensions=(0,)
+    "test_blackjax_method"
+    "test_legacy_nuts_blackjax_warning"
+
     # Tests require network access
     "test_alias_equal_to_name"
     "test_average_by"