From f031c160d692f4773439c7b8b7e1586645ca9295 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Tue, 7 Nov 2023 15:22:29 +0100
Subject: [PATCH 1/2] python311Packages.flax: 0.7.4 -> 0.7.5

Changelog: https://github.com/google/flax/releases/tag/v0.7.5
---
 .../python-modules/flax/default.nix           | 29 ++++++++++++-------
 1 file changed, 18 insertions(+), 11 deletions(-)

diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix
index 0a6b086fe516..b68f56a11f41 100644
--- a/pkgs/development/python-modules/flax/default.nix
+++ b/pkgs/development/python-modules/flax/default.nix
@@ -1,36 +1,43 @@
-{ buildPythonPackage
+{ lib
+, buildPythonPackage
 , fetchFromGitHub
 , jaxlib
+, pythonRelaxDepsHook
+, setuptools-scm
+, cloudpickle
 , jax
-, keras
-, lib
 , matplotlib
 , msgpack
 , numpy
 , optax
+, rich
+, tensorstore
+, keras
 , pytest-xdist
 , pytestCheckHook
-, pythonRelaxDepsHook
 , tensorflow
-, tensorstore
-, fetchpatch
-, rich
 }:
 
 buildPythonPackage rec {
   pname = "flax";
-  version = "0.7.4";
+  version = "0.7.5";
+  pyproject = true;
 
   src = fetchFromGitHub {
     owner = "google";
-    repo = pname;
+    repo = "flax";
     rev = "refs/tags/v${version}";
-    hash = "sha256-i48omag/1Si3mCCGfsUD9qeejyeCLWzvvwKJqH8vm8k=";
+    hash = "sha256-NDah0ayQbiO1/sTU1DDf/crPq5oLTnSuosV7cFHlTM8=";
   };
 
-  nativeBuildInputs = [ jaxlib pythonRelaxDepsHook ];
+  nativeBuildInputs = [
+    jaxlib
+    pythonRelaxDepsHook
+    setuptools-scm
+  ];
 
   propagatedBuildInputs = [
+    cloudpickle
     jax
     matplotlib
     msgpack

From ec24f13c269d9b00f1758094888f8fe19f362ff1 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Tue, 7 Nov 2023 15:35:12 +0100
Subject: [PATCH 2/2] python311Packages.dalle-mini: mark as broken

---
 pkgs/development/python-modules/dalle-mini/default.nix | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pkgs/development/python-modules/dalle-mini/default.nix b/pkgs/development/python-modules/dalle-mini/default.nix
index f20b693d3870..e50249dc7dd9 100644
--- a/pkgs/development/python-modules/dalle-mini/default.nix
+++ b/pkgs/development/python-modules/dalle-mini/default.nix
@@ -49,5 +49,7 @@ buildPythonPackage rec {
     homepage = "https://github.com/borisdayma/dalle-mini";
     license = licenses.asl20;
     maintainers = with maintainers; [ r-burns ];
+    # incompatible with recent versions of JAX
+    broken = true;
   };
 }