From 2dbcd73818222ca6c53093c18c1fac232847fc93 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Thu, 17 Oct 2024 08:55:34 +0200
Subject: [PATCH] python312Packages.flax: 0.9.0 -> 0.10.1

Diff: https://github.com/google/flax/compare/refs/tags/v0.9.0...v0.10.1

Changelog: https://github.com/google/flax/releases/tag/v0.10.1
---
 pkgs/development/python-modules/flax/default.nix | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix
index 4f93bd4f8ea5..92f95f3aef7d 100644
--- a/pkgs/development/python-modules/flax/default.nix
+++ b/pkgs/development/python-modules/flax/default.nix
@@ -21,9 +21,11 @@
   # checks
   cloudpickle,
   einops,
+  flaxlib,
   keras,
-  pytest-xdist,
   pytestCheckHook,
+  pytest-xdist,
+  sphinx,
   tensorflow,
   treescope,
 
@@ -33,14 +35,14 @@
 
 buildPythonPackage rec {
   pname = "flax";
-  version = "0.9.0";
+  version = "0.10.1";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "google";
     repo = "flax";
     rev = "refs/tags/v${version}";
-    hash = "sha256-iDWuUJKO7V4QrbVsS4ALgy6fbllOC43o7W4mhjtZ9xc=";
+    hash = "sha256-+URbQGnmqmSNgucEyWvI5DMnzXjpmJzLA+Pho2lX+S4=";
   };
 
   build-system = [
@@ -69,9 +71,11 @@ buildPythonPackage rec {
   nativeCheckInputs = [
     cloudpickle
     einops
+    flaxlib
     keras
-    pytest-xdist
     pytestCheckHook
+    pytest-xdist
+    sphinx
     tensorflow
     treescope
   ];