python: Add python-jaxopt.

* guix-science/packages/python.scm (python-jaxopt): New variable.

1 files changed, 56 insertions, 0 deletions

diff --git a/guix-science/packages/python.scm b/guix-science/packages/python.scm
index 778c9a7..cfb13b3 100644
--- a/guix-science/packages/python.scm
+++ b/guix-science/packages/python.scm
@@ -769,6 +769,62 @@ as forward-mode differentiation, and the two can be composed
 arbitrarily to any order.")
     (license license:asl2.0)))
 
+(define-public python-jaxopt
+  (package
+    (name "python-jaxopt")
+    (version "0.8.3")
+    (source
+     (origin
+       (method git-fetch)
+       (uri (git-reference
+             (url "https://github.com/google/jaxopt")
+             (commit (string-append "jaxopt-v" version))))
+       (file-name (git-file-name name version))
+       (sha256
+        (base32 "0mak1rpjzw6v8birinrfn1plilgyvdxxqqs05rp89p54gd54gw2g"))))
+    (build-system pyproject-build-system)
+    (arguments
+     (list #:test-flags
+           ;; Differences in output due to strict tolerances.
+           '(list "-k" (string-append
+                        "not test_binary_logit_log_likelihood"
+                        " and not test_compare_with_sklearn0"
+                        " and not test_compare_with_sklearn1"
+                        " and not test_against_scipy0"
+                        " and not test_against_scipy1"))))
+    (propagated-inputs
+     (list python-dm-tree
+           python-jax
+           python-jaxlib
+           python-optax
+           python-numpy
+           python-scipy))
+    (native-inputs
+     (list python-absl-py
+           python-cvxopt
+           python-cvxpy
+           python-pytest
+           python-pytest-xdist
+           python-scikit-learn
+           python-setuptools
+           python-wheel))
+    (home-page "https://github.com/google/jaxopt")
+    (synopsis "Hardware accelerated, batchable and differentiable optimizers in JAX")
+    (description "JAXopt provides hardware accelerated, batchable and
+differentiable optimizers in JAX.
+
+@enumerate
+@item Hardware accelerated: the implementations run on GPU and TPU, in
+  addition to CPU.
+@item Batchable: multiple instances of the same optimization problem
+  can be automatically vectorized using JAX’s @code{vmap}.
+@item Differentiable: optimization problem solutions can be
+  differentiated with respect to their inputs either implicitly or via
+  autodiff of unrolled algorithm iterations.
+@end enumerate
+")
+    (license license:asl2.0)))
+
 (define tensorflow-system-libs
   (list
    ;;"absl_py"