From 03cfe274f88306765d1518115450458c3f41c536 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Mon, 1 Apr 2024 14:09:44 +0200
Subject: [PATCH] python311Packages.optax: 0.2.1 -> 0.2.2

Changelog: https://github.com/google-deepmind/optax/releases/tag/v0.2.2
---
 pkgs/development/python-modules/optax/default.nix | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/pkgs/development/python-modules/optax/default.nix b/pkgs/development/python-modules/optax/default.nix
index dd36a90864f2..950e9c91dbf5 100644
--- a/pkgs/development/python-modules/optax/default.nix
+++ b/pkgs/development/python-modules/optax/default.nix
@@ -1,6 +1,7 @@
 { lib
 , absl-py
 , buildPythonPackage
+, flit-core
 , chex
 , fetchFromGitHub
 , jaxlib
@@ -11,16 +12,16 @@
 
 buildPythonPackage rec {
   pname = "optax";
-  version = "0.2.1";
-  format = "setuptools";
+  version = "0.2.2";
+  pyproject = true;
 
-  disabled = pythonOlder "3.7";
+  disabled = pythonOlder "3.9";
 
   src = fetchFromGitHub {
     owner = "deepmind";
-    repo = pname;
+    repo = "optax";
     rev = "refs/tags/v${version}";
-    hash = "sha256-vimsVZV5Z11euLxsu998pMQZ0hG3xl96D3h9iONtl/E=";
+    hash = "sha256-sBiKUuQR89mttc9Njrh1aeUJOYdlcF7Nlj3/+Y7OMb4=";
   };
 
   outputs = [
@@ -28,6 +29,10 @@ buildPythonPackage rec {
     "testsout"
   ];
 
+  nativeBuildInputs = [
+    flit-core
+  ];
+
   buildInputs = [
     jaxlib
   ];