From 148c0ec36ef03d2c9d12676fa4560d09821d01ea Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 18 Sep 2024 07:53:13 +0200
Subject: [PATCH 1/3] python312Packages.numpyro: 0.15.2 -> 0.15.3

Diff: https://github.com/pyro-ppl/numpyro/compare/0.15.2...0.15.3

Changelog: https://github.com/pyro-ppl/numpyro/releases/tag/0.15.3
---
 pkgs/development/python-modules/numpyro/default.nix | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/pkgs/development/python-modules/numpyro/default.nix b/pkgs/development/python-modules/numpyro/default.nix
index dd5aca045f43..945a33ef0053 100644
--- a/pkgs/development/python-modules/numpyro/default.nix
+++ b/pkgs/development/python-modules/numpyro/default.nix
@@ -1,7 +1,6 @@
 {
   lib,
   buildPythonPackage,
-  pythonOlder,
   fetchPypi,
 
   # build-system
@@ -23,14 +22,12 @@
 
 buildPythonPackage rec {
   pname = "numpyro";
-  version = "0.15.2";
+  version = "0.15.3";
   pyproject = true;
 
-  disabled = pythonOlder "3.9";
-
   src = fetchPypi {
     inherit version pname;
-    hash = "sha256-6G3TrDyQ5N2uuzLzzEus1czCtvg3M0wBorLo2vQZozE=";
+    hash = "sha256-9EXurmIA+IPXkNZc4p/yRfepY5usMyLZk+zPkcRAI7M=";
   };
 
   build-system = [ setuptools ];

From 5ca6aa4004d944387e7580060342d4954e1bfddf Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 18 Sep 2024 08:58:33 +0200
Subject: [PATCH 2/3] python312Packages.numpyro: switch to fetchFromGitHub

---
 .../python-modules/numpyro/default.nix        | 63 ++++++++++++++++---
 1 file changed, 53 insertions(+), 10 deletions(-)

diff --git a/pkgs/development/python-modules/numpyro/default.nix b/pkgs/development/python-modules/numpyro/default.nix
index 945a33ef0053..1df3800c4482 100644
--- a/pkgs/development/python-modules/numpyro/default.nix
+++ b/pkgs/development/python-modules/numpyro/default.nix
@@ -1,7 +1,7 @@
 {
   lib,
   buildPythonPackage,
-  fetchPypi,
+  fetchFromGitHub,
 
   # build-system
   setuptools,
@@ -13,11 +13,19 @@
   numpy,
   tqdm,
 
-  # checks
+  # tests
+  # Our current version of tensorflow (2.13.0) is too old and doesn't support python>=3.12
+  # We remove optional test dependencies that require tensorflow and skip the corresponding tests to
+  # avoid introducing a useless incompatibility with python 3.12:
+  # dm-haiku,
+  # flax,
+  # tensorflow-probability,
   funsor,
+  graphviz,
+  optax,
+  pyro-api,
   pytestCheckHook,
-# TODO: uncomment when tensorflow-probability gets fixed.
-# , tensorflow-probability
+  scikit-learn,
 }:
 
 buildPythonPackage rec {
@@ -25,9 +33,11 @@ buildPythonPackage rec {
   version = "0.15.3";
   pyproject = true;
 
-  src = fetchPypi {
-    inherit version pname;
-    hash = "sha256-9EXurmIA+IPXkNZc4p/yRfepY5usMyLZk+zPkcRAI7M=";
+  src = fetchFromGitHub {
+    owner = "pyro-ppl";
+    repo = "numpyro";
+    rev = "refs/tags/${version}";
+    hash = "sha256-g+ep221hhLbCjQasKpiEAXkygI5A3Hglqo1tV8lv5eg=";
   };
 
   build-system = [ setuptools ];
@@ -41,9 +51,14 @@ buildPythonPackage rec {
   ];
 
   nativeCheckInputs = [
+    # dm-haiku
+    # flax
     funsor
+    graphviz
+    optax
+    pyro-api
     pytestCheckHook
-    # TODO: uncomment when tensorflow-probability gets fixed.
+    scikit-learn
     # tensorflow-probability
   ];
 
@@ -60,17 +75,45 @@ buildPythonPackage rec {
     "test_kl_dirichlet_dirichlet"
     "test_kl_univariate"
     "test_mean_var"
+
     # Tests want to download data
     "data_load"
     "test_jsb_chorales"
+
     # RuntimeWarning: overflow encountered in cast
     "test_zero_inflated_logits_probs_agree"
+
     # NameError: unbound axis name: _provenance
     "test_model_transformation"
+
+    # require dm-haiku
+    "test_flax_state_dropout_smoke"
+    "test_flax_module"
+    "test_random_module_mcmc"
+
+    # require flax
+    "test_haiku_state_dropout_smoke"
+    "test_haiku_module"
+    "test_random_module_mcmc"
+
+    # require tensorflow-probability
+    "test_modified_bessel_first_kind_vect"
+    "test_diag_spectral_density_periodic"
+    "test_kernel_approx_periodic"
+    "test_modified_bessel_first_kind_one_dim"
+    "test_modified_bessel_first_kind_vect"
+    "test_periodic_gp_one_dim_model"
+    "test_no_tracer_leak_at_lazy_property_sample"
   ];
 
-  # TODO: remove when tensorflow-probability gets fixed.
-  disabledTestPaths = [ "test/test_distributions.py" ];
+  disabledTestPaths = [
+    # require jaxns (unpackaged)
+    "test/contrib/test_nested_sampling.py"
+
+    # requires tensorflow-probability
+    "test/contrib/test_tfp.py"
+    "test/test_distributions.py"
+  ];
 
   meta = {
     description = "Library for probabilistic programming with NumPy";

From 33278de087b75243bd6e7480bcd622a2c5279201 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Sat, 21 Sep 2024 18:34:51 +0200
Subject: [PATCH 3/3] python312Packages.numpyro: disable failing test on darwin

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

diff --git a/pkgs/development/python-modules/numpyro/default.nix b/pkgs/development/python-modules/numpyro/default.nix
index 1df3800c4482..a1aed20a5268 100644
--- a/pkgs/development/python-modules/numpyro/default.nix
+++ b/pkgs/development/python-modules/numpyro/default.nix
@@ -104,6 +104,11 @@ buildPythonPackage rec {
     "test_modified_bessel_first_kind_vect"
     "test_periodic_gp_one_dim_model"
     "test_no_tracer_leak_at_lazy_property_sample"
+
+    # flaky on darwin
+    # TODO: uncomment at next release (0.15.4) as it has been fixed:
+    # https://github.com/pyro-ppl/numpyro/pull/1863
+    "test_change_point_x64"
   ];
 
   disabledTestPaths = [