From 54fd3e5b7f3eeea588ecd2d1c4a4c7bc14fab6a9 Mon Sep 17 00:00:00 2001
From: Jonas Heinrich <onny@project-insanity.org>
Date: Thu, 18 Aug 2022 13:23:07 +0200
Subject: [PATCH] python310Packages.flax: 0.5.2 -> 0.6.0

---
 .../python-modules/flax/default.nix           | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix
index 5eda127fb1c0..ebdb3ac898ad 100644
--- a/pkgs/development/python-modules/flax/default.nix
+++ b/pkgs/development/python-modules/flax/default.nix
@@ -1,6 +1,7 @@
 { buildPythonPackage
 , fetchFromGitHub
 , jaxlib
+, jax
 , keras
 , lib
 , matplotlib
@@ -10,26 +11,40 @@
 , pytest-xdist
 , pytestCheckHook
 , tensorflow
+, fetchpatch
+, rich
 }:
 
 buildPythonPackage rec {
   pname = "flax";
-  version = "0.5.2";
+  version = "0.6.0";
 
   src = fetchFromGitHub {
     owner = "google";
     repo = pname;
     rev = "refs/tags/v${version}";
-    sha256 = "sha256-t24JZ08EmvuAINZC26OQI1icklUhUkfz6ZRKPr2COAw=";
+    sha256 = "sha256-egTYYFZxhE/Kk7jXRi1HmjCjyFia2LoRigH042isDu0=";
   };
 
+  patches = [
+    # Bump rich dependency, should be fixed in releases after 0.6.0
+    # https://github.com/google/flax/pull/2407
+    (fetchpatch {
+      url = "https://github.com/google/flax/commit/72189153f9779022b97858ae747c23fbaf571e3d.patch";
+      sha256 = "sha256-hKOn/M7qpBM6R1RIJpnXpRoZgIHqkwQZApN4L0fBzIE=";
+      name = "bump_rich_dependency.patch";
+    })
+  ];
+
   buildInputs = [ jaxlib ];
 
   propagatedBuildInputs = [
+    jax
     matplotlib
     msgpack
     numpy
     optax
+    rich
   ];
 
   pythonImportsCheck = [