zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-30 07:18:15 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

non-symmetric eig and eigh (#2188)

Browse Source
This commit is contained in:
Awni Hannun
2025-05-15 13:01:44 -07:00
committed by GitHub
parent cf6c939e86
commit c1eb9d05d9
14 changed files with 423 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

72
python/src/linalg.cpp
View File

@@ -236,7 +236,7 @@ void init_linalg(nb::module_& parent_module) {
Returns:
Union[tuple(array, ...), array]:
If compute_uv is ``True`` returns the ``U``, ``S``, and ``Vt`` matrices, such that
If compute_uv is ``True`` returns the ``U``, ``S``, and ``Vt`` matrices, such that
``A = U @ diag(S) @ Vt``. If compute_uv is ``False`` returns singular values array ``S``.
)pbdoc");
m.def(
@@ -407,6 +407,76 @@ void init_linalg(nb::module_& parent_module) {
Returns:
array: The cross product of ``a`` and ``b`` along the specified axis.
)pbdoc");
m.def(
"eigvals",
&mx::linalg::eigvals,
"a"_a,
nb::kw_only(),
"stream"_a = nb::none(),
R"pbdoc(
Compute the eigenvalues of a square matrix.
This function differs from :func:`numpy.linalg.eigvals` in that the
return type is always complex even if the eigenvalues are all real.
This function supports arrays with at least 2 dimensions. When the
input has more than two dimensions, the eigenvalues are computed for
each matrix in the last two dimensions.
Args:
a (array): The input array.
stream (Stream, optional): Stream or device. Defaults to ``None``
in which case the default stream of the default device is used.
Returns:
array: The eigenvalues (not necessarily in order).
Example:
>>> A = mx.array([[1., -2.], [-2., 1.]])
>>> eigenvalues = mx.linalg.eigvals(A, stream=mx.cpu)
>>> eigenvalues
array([3+0j, -1+0j], dtype=complex64)
)pbdoc");
m.def(
"eig",
[](const mx::array& a, mx::StreamOrDevice s) {
auto result = mx::linalg::eig(a, s);
return nb::make_tuple(result.first, result.second);
},
"a"_a,
nb::kw_only(),
"stream"_a = nb::none(),
R"pbdoc(
Compute the eigenvalues and eigenvectors of a square matrix.
This function differs from :func:`numpy.linalg.eig` in that the
return type is always complex even if the eigenvalues are all real.
This function supports arrays with at least 2 dimensions. When the input
has more than two dimensions, the eigenvalues and eigenvectors are
computed for each matrix in the last two dimensions.
Args:
a (array): The input array.
stream (Stream, optional): Stream or device. Defaults to ``None``
in which case the default stream of the default device is used.
Returns:
Tuple[array, array]:
A tuple containing the eigenvalues and the normalized right
eigenvectors. The column ``v[:, i]`` is the eigenvector
corresponding to the i-th eigenvalue.
Example:
>>> A = mx.array([[1., -2.], [-2., 1.]])
>>> w, v = mx.linalg.eig(A, stream=mx.cpu)
>>> w
array([3+0j, -1+0j], dtype=complex64)
>>> v
array([[0.707107+0j, 0.707107+0j],
[-0.707107+0j, 0.707107+0j]], dtype=complex64)
)pbdoc");
m.def(
"eigvalsh",
&mx::linalg::eigvalsh,

77
python/tests/test_linalg.py
View File

@@ -312,6 +312,83 @@ class TestLinalg(mlx_tests.MLXTestCase):
with self.assertRaises(ValueError):
mx.linalg.cross(a, b)
def test_eig(self):
tols = {"atol": 1e-5, "rtol": 1e-5}
def check_eigs_and_vecs(A_np, kwargs={}):
A = mx.array(A_np)
eig_vals, eig_vecs = mx.linalg.eig(A, stream=mx.cpu, **kwargs)
self.assertTrue(
mx.allclose(A @ eig_vecs, eig_vals[..., None, :] * eig_vecs, **tols)
)
eig_vals_only = mx.linalg.eigvals(A, stream=mx.cpu, **kwargs)
self.assertTrue(mx.allclose(eig_vals, eig_vals_only, **tols))
# Test a simple 2x2 matrix
A_np = np.array([[1.0, 1.0], [3.0, 4.0]], dtype=np.float32)
check_eigs_and_vecs(A_np)
# Test complex eigenvalues
A_np = np.array([[1.0, -1.0], [1.0, 1.0]], dtype=np.float32)
check_eigs_and_vecs(A_np)
# Test a larger random symmetric matrix
n = 5
np.random.seed(1)
A_np = np.random.randn(n, n).astype(np.float32)
check_eigs_and_vecs(A_np)
# Test with batched input
A_np = np.random.randn(3, n, n).astype(np.float32)
check_eigs_and_vecs(A_np)
# Test error cases
with self.assertRaises(ValueError):
mx.linalg.eig(mx.array([1.0, 2.0])) # 1D array
with self.assertRaises(ValueError):
mx.linalg.eig(
mx.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
) # Non-square matrix
with self.assertRaises(ValueError):
mx.linalg.eigvals(mx.array([1.0, 2.0])) # 1D array
with self.assertRaises(ValueError):
mx.linalg.eigvals(
mx.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
) # Non-square matrix
def test_lu(self):
with self.assertRaises(ValueError):
mx.linalg.lu(mx.array(0.0), stream=mx.cpu)
with self.assertRaises(ValueError):
mx.linalg.lu(mx.array([0.0, 1.0]), stream=mx.cpu)
with self.assertRaises(ValueError):
mx.linalg.lu(mx.array([[0, 1], [1, 0]]), stream=mx.cpu)
# Test 3x3 matrix
a = mx.array([[3.0, 1.0, 2.0], [1.0, 8.0, 6.0], [9.0, 2.0, 5.0]])
P, L, U = mx.linalg.lu(a, stream=mx.cpu)
self.assertTrue(mx.allclose(L[P, :] @ U, a))
# Test batch dimension
a = mx.broadcast_to(a, (5, 5, 3, 3))
P, L, U = mx.linalg.lu(a, stream=mx.cpu)
L = mx.take_along_axis(L, P[..., None], axis=-2)
self.assertTrue(mx.allclose(L @ U, a))
# Test non-square matrix
a = mx.array([[3.0, 1.0, 2.0], [1.0, 8.0, 6.0]])
P, L, U = mx.linalg.lu(a, stream=mx.cpu)
self.assertTrue(mx.allclose(L[P, :] @ U, a))
a = mx.array([[3.0, 1.0], [1.0, 8.0], [9.0, 2.0]])
P, L, U = mx.linalg.lu(a, stream=mx.cpu)
self.assertTrue(mx.allclose(L[P, :] @ U, a))
def test_eigh(self):
tols = {"atol": 1e-5, "rtol": 1e-5}