python3Packages.blackjax: 1.3 -> 1.4 (#503850)

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"

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"