From 25ce86cd74f38d7c17d1d90a8271758959702f43 Mon Sep 17 00:00:00 2001
From: Ricardo Wurmus <rekado@elephly.net>
Date: Mon, 30 Oct 2023 15:40:49 +0100
Subject: python: Add python-optax.

* guix-science/packages/python.scm (python-optax): New variable.
---
 guix-science/packages/python.scm | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/guix-science/packages/python.scm b/guix-science/packages/python.scm
index cc85c02..df03161 100644
--- a/guix-science/packages/python.scm
+++ b/guix-science/packages/python.scm
@@ -1997,3 +1997,33 @@ reliable JAX code.  This includes utils to help:
 @end itemize
 ")
     (license license:asl2.0)))
+
+(define-public python-optax
+  (package
+    (name "python-optax")
+    ;; 0.1.6 needs a more recent numpy
+    (version "0.1.5")
+    (source
+     (origin
+       (method url-fetch)
+       (uri (pypi-uri "optax" version))
+       (sha256
+        (base32 "0bhgaaxvqli3b2081zp9ycb2c2hqba0fcpigaqjxbnsidyspk8qa"))))
+    (build-system pyproject-build-system)
+    ;; Tests require haiku, tensorflow, and flax, but flax needs
+    ;; optax.
+    (arguments
+     (list #:tests? #false))
+    (propagated-inputs (list python-absl-py
+                             python-chex
+                             python-jax
+                             python-jaxlib
+                             python-numpy))
+    (native-inputs
+     (list python-dm-tree
+           python-pytest))
+    (home-page "https://github.com/google-deepmind/optax/")
+    (synopsis "Gradient processing and optimization library for JAX")
+    (description "Optax is a gradient processing and optimization
+library for JAX.")
+    (license license:asl2.0)))
-- 
cgit v1.2.3