From 20f49fa10b8d28c581a1cf125de12e5bbf47840a Mon Sep 17 00:00:00 2001
From: Ben Darwin <bcdarwin@gmail.com>
Date: Mon, 6 Feb 2023 21:54:43 -0500
Subject: [PATCH] python310Packages.blackjax: init at 0.9.6

---
 .../python-modules/blackjax/default.nix       | 62 +++++++++++++++++++
 pkgs/top-level/python-packages.nix            |  2 +
 2 files changed, 64 insertions(+)
 create mode 100644 pkgs/development/python-modules/blackjax/default.nix

diff --git a/pkgs/development/python-modules/blackjax/default.nix b/pkgs/development/python-modules/blackjax/default.nix
new file mode 100644
index 000000000000..62fdc0e94a98
--- /dev/null
+++ b/pkgs/development/python-modules/blackjax/default.nix
@@ -0,0 +1,62 @@
+{ lib
+, buildPythonPackage
+, pythonOlder
+, fetchFromGitHub
+, fetchpatch
+, pytestCheckHook
+, fastprogress
+, jax
+, jaxlib
+, jaxopt
+, optax
+, typing-extensions
+}:
+
+buildPythonPackage rec {
+  pname = "blackjax";
+  version = "0.9.6";
+  disabled = pythonOlder "3.7";
+
+  src = fetchFromGitHub {
+    owner = "blackjax-devs";
+    repo = pname;
+    rev = "refs/tags/${version}";
+    hash = "sha256-EieDu9SJxi2cp1bHlxX4vvFZeDGMGIm24GoR8nSyjvE=";
+  };
+
+  patches = [
+    # remove in next release
+    (fetchpatch {
+      name = "fix-lbfgs-args";
+      url = "https://github.com/blackjax-devs/blackjax/commit/1aaa6f64bbcb0557b658604b2daba826e260cbc6.patch";
+      hash = "sha256-XyjorXPH5Ap35Tv1/lTeTWamjplJF29SsvOq59ypftE=";
+    })
+  ];
+
+  propagatedBuildInputs = [
+    fastprogress
+    jax
+    jaxlib
+    jaxopt
+    optax
+    typing-extensions
+  ];
+
+  nativeCheckInputs = [ pytestCheckHook ];
+  disabledTestPaths = [ "tests/test_benchmarks.py" ];
+  disabledTests = [
+    # too slow
+    "test_adaptive_tempered_smc"
+  ];
+
+  pythonImportsCheck = [
+    "blackjax"
+  ];
+
+  meta = with lib; {
+    homepage = "https://blackjax-devs.github.io/blackjax";
+    description = "Sampling library designed for ease of use, speed and modularity";
+    license = licenses.asl20;
+    maintainers = with maintainers; [ bcdarwin ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 3d1502c30ddf..71b779dfab40 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -1296,6 +1296,8 @@ self: super: with self; {
 
   black = callPackage ../development/python-modules/black { };
 
+  blackjax = callPackage ../development/python-modules/blackjax { };
+
   black-macchiato = callPackage ../development/python-modules/black-macchiato { };
 
   bleach = callPackage ../development/python-modules/bleach { };