Merge pull request data-apis#379 from ev-br/add_eig

ENH: add compatibility shims for {eig,eigvals}

Author: ev-br

array_api_compat/common/_linalg.py

@@ -34,6 +34,10 @@ class EighResult(NamedTuple):
     eigenvalues: Array
     eigenvectors: Array
 
+class EigResult(NamedTuple):
+    eigenvalues: Array
+    eigenvectors: Array
+
 class QRResult(NamedTuple):
     Q: Array
     R: Array

array_api_compat/numpy/linalg.py

@@ -19,6 +19,7 @@
 cross = get_xp(np)(_linalg.cross)
 outer = get_xp(np)(_linalg.outer)
 EighResult = _linalg.EighResult
+EigResult = _linalg.EigResult
 QRResult = _linalg.QRResult
 SlogdetResult = _linalg.SlogdetResult
 SVDResult = _linalg.SVDResult
@@ -97,6 +98,85 @@ def solve(x1: Array, x2: Array, /) -> Array:
     return wrap(r.astype(result_t, copy=False))
 
 
+# Unlike numpy.linalg.eig, Array API version always returns complex results
+
+def eig(x: Array, /) -> tuple[Array, Array]:
+    try:
+        from numpy.linalg._linalg import (  # type: ignore[attr-defined]
+            _assert_stacked_square,
+            _assert_finite,
+            _commonType,
+            _makearray,
+            _raise_linalgerror_eigenvalues_nonconvergence,
+            isComplexType,
+            _complexType,
+        )
+    except ImportError:
+        from numpy.linalg.linalg import (  # type: ignore[attr-defined]
+            _assert_stacked_square,
+            _assert_finite,
+            _commonType,
+            _makearray,
+            _raise_linalgerror_eigenvalues_nonconvergence,
+            isComplexType,
+            _complexType,
+        )
+    from numpy.linalg import _umath_linalg
+
+    x, wrap = _makearray(x)
+    _assert_stacked_square(x)
+    _assert_finite(x)
+    t, result_t = _commonType(x)
+
+    signature = 'D->DD' if isComplexType(t) else 'd->DD'
+    with np.errstate(call=_raise_linalgerror_eigenvalues_nonconvergence,
+                  invalid='call', over='ignore', divide='ignore',
+                  under='ignore'):
+        w, vt = _umath_linalg.eig(x, signature=signature)
+
+    result_t = _complexType(result_t)
+    vt = vt.astype(result_t, copy=False)
+    return EigResult(w.astype(result_t, copy=False), wrap(vt))
+
+
+def eigvals(x: Array, /) -> Array:
+    try:
+        from numpy.linalg._linalg import (  # type: ignore[attr-defined]
+            _assert_stacked_square,
+            _assert_finite,
+            _commonType,
+            _makearray,
+            _raise_linalgerror_eigenvalues_nonconvergence,
+            isComplexType,
+            _complexType,
+        )
+    except ImportError:
+        from numpy.linalg.linalg import (  # type: ignore[attr-defined]
+            _assert_stacked_square,
+            _assert_finite,
+            _commonType,
+            _makearray,
+            _raise_linalgerror_eigenvalues_nonconvergence,
+            isComplexType,
+            _complexType,
+        )
+    from numpy.linalg import _umath_linalg
+
+    x, wrap = _makearray(x)
+    _assert_stacked_square(x)
+    _assert_finite(x)
+    t, result_t = _commonType(x)
+
+    signature = 'D->D' if isComplexType(t) else 'd->D'
+    with np.errstate(call=_raise_linalgerror_eigenvalues_nonconvergence,
+                  invalid='call', over='ignore', divide='ignore',
+                  under='ignore'):
+        w = _umath_linalg.eigvals(x, signature=signature)
+
+    result_t = _complexType(result_t)
+    return w.astype(result_t, copy=False)
+
+
 # These functions are completely new here. If the library already has them
 # (i.e., numpy 2.0), use the library version instead of our wrapper.
 if hasattr(np.linalg, "vector_norm"):

cupy-xfails.txt

@@ -24,6 +24,10 @@ array_api_tests/test_has_names.py::test_has_names[array_attribute-mT]
 
 array_api_tests/test_linalg.py::test_solve
 
+# 2025.12 support; {eig,eigvals} are new in CuPy 14
+array_api_tests/test_linalg.py::test_eig
+array_api_tests/test_linalg.py::test_eigvals
+
 # We cannot modify array methods
 array_api_tests/test_operators_and_elementwise_functions.py::test_divide[__truediv__(x, s)]
 array_api_tests/test_operators_and_elementwise_functions.py::test_floor_divide[__floordiv__(x, s)]

dask-xfails.txt

@@ -136,6 +136,9 @@ array_api_tests/test_data_type_functions.py::TestBroadcastShapes::test_broadcast
 array_api_tests/test_data_type_functions.py::TestBroadcastShapes::test_empty
 array_api_tests/test_data_type_functions.py::TestBroadcastShapes::test_error
 
+array_api_tests/test_linalg.py::test_eig
+array_api_tests/test_linalg.py::test_eigvals
+
 # Stubs have a comment: (**note**: libraries may return ``NaN`` to match Python behavior.)
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]