From 14e0bca6ddb2f30e943962d6bee50a6da5df4c67 Mon Sep 17 00:00:00 2001
From: Nicolas Graves <ngraves@ngraves.fr>
Date: Sat, 30 Aug 2025 03:20:18 +0200
Subject: [PATCH] gnu: Add python-ml-dtypes.

* gnu/packages/machine-learning.scm (python-ml-dtypes): New variable.

Change-Id: I127edac9317d6380d22c41b6c6d578e183ab6b9f
Signed-off-by: Andreas Enge <andreas@enge.fr>
---
 gnu/packages/machine-learning.scm | 59 +++++++++++++++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/gnu/packages/machine-learning.scm b/gnu/packages/machine-learning.scm
index 4c3f20c669e06e769f157ae1a9d856bc4acc0a06..82fa0cd5745fa0df628cd2c813761fd6365daeb9 100644
--- a/gnu/packages/machine-learning.scm
+++ b/gnu/packages/machine-learning.scm
@@ -519,6 +519,65 @@ transforms.")
 Learning usecases.")
     (license license:asl2.0)))
 
+(define-public python-ml-dtypes
+  (package
+    (name "python-ml-dtypes")
+    (version "0.2.0")
+    (source
+     (origin
+       (method git-fetch)
+       (uri (git-reference
+              (url "https://github.com/jax-ml/ml_dtypes")
+              (commit (string-append "v" version))))
+       (file-name (git-file-name name version))
+       (sha256
+        (base32 "1yv90f28c9w34430xjwvn1lzxdylvp1zi6b02cx7crla6qkvrzn5"))))
+    (build-system pyproject-build-system)
+    (arguments
+     (list
+      #:phases
+      #~(modify-phases %standard-phases
+          (add-after 'unpack 'use-eigen-package
+            (lambda* (#:key inputs #:allow-other-keys)
+              (substitute* "setup.py"
+                (("third_party/eigen")
+                 (search-input-directory inputs "include/eigen3")))))
+          (add-after 'install 'symlink-lib
+            (lambda* (#:key inputs outputs #:allow-other-keys)
+              (for-each
+               (lambda (file)
+                 (symlink file
+                          (string-append #$output "/lib/ml_dtypes.so")))
+               (find-files (site-packages inputs outputs) "\\.so$")))))))
+    (inputs (list eigen-for-python-ml-dtypes))
+    (propagated-inputs (list python-numpy))
+    (native-inputs (list pybind11
+                         python-absl-py
+                         python-pytest
+                         python-setuptools
+                         python-wheel))
+    (home-page "https://github.com/jax-ml/ml_dtypes")
+    (synopsis "NumPy dtype extensions used in machine learning")
+    (description
+     "This package is a stand-alone implementation of several
+NumPy @code{dtype} extensions used in machine learning libraries, including:
+
+@itemize
+@item @code{bfloat16}: an alternative to the standard @code{float16} format
+@item @code{float8_*}: several experimental 8-bit floating point
+  representations including:
+  @itemize
+  @item @code{float8_e4m3b11fnuz}
+  @item @code{float8_e4m3fn}
+  @item @code{float8_e4m3fnuz}
+  @item @code{float8_e5m2}
+  @item @code{float8_e5m2fnuz}
+  @end itemize
+@item @code{int4} and @code{uint4}: low precision integer types.
+@end itemize
+")
+    (license license:asl2.0)))
+
 (define-public ghmm
   ;; The latest release candidate is several years and a couple of fixes have
   ;; been published since.  This is why we download the sources from the SVN