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

Eigenvalues and eigenvectors (#1334)

Browse Source 
* initial eigvalsh

* add compute_vectors

* add compute_vectors_

* return a pair

* add eigh to return only eigenvectors

* fixed typo

* merge merge Eighvalsh and Eigh into a single primitive

* use the same primate with the flag

* fix primatives

* use MULTI

* fix eval_gpu

* fix decleration

* rename EighPrimitive to Eigh

* tests

* tests

* fix rebase and format

* cleanup lapack

* format

* add cblas.h

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Kashif Rasul
2024-10-22 21:18:48 +02:00
committed by GitHub
parent c26208f67d
commit 3ddc07e936
23 changed files with 434 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
python/tests/test_linalg.py
View File

@@ -268,6 +268,57 @@ class TestLinalg(mlx_tests.MLXTestCase):
with self.assertRaises(ValueError):
mx.linalg.cross(a, b)
def test_eigh(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.eigh(A, stream=mx.cpu, **kwargs)
eig_vals_np, _ = np.linalg.eigh(A_np, **kwargs)
self.assertTrue(np.allclose(eig_vals, eig_vals_np, **tols))
self.assertTrue(
mx.allclose(A @ eig_vecs, eig_vals[..., None, :] * eig_vecs, **tols)
)
eig_vals_only = mx.linalg.eigvalsh(A, stream=mx.cpu, **kwargs)
self.assertTrue(mx.allclose(eig_vals, eig_vals_only, **tols))
# Test a simple 2x2 symmetric matrix
A_np = np.array([[1.0, 2.0], [2.0, 4.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)
A_np = (A_np + A_np.T) / 2
check_eigs_and_vecs(A_np)
# Test with upper triangle
check_eigs_and_vecs(A_np, {"UPLO": "U"})
# Test with batched input
A_np = np.random.randn(3, n, n).astype(np.float32)
A_np = (A_np + np.transpose(A_np, (0, 2, 1))) / 2
check_eigs_and_vecs(A_np)
# Test error cases
with self.assertRaises(ValueError):
mx.linalg.eigh(mx.array([1.0, 2.0])) # 1D array
with self.assertRaises(ValueError):
mx.linalg.eigh(
mx.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
) # Non-square matrix
with self.assertRaises(ValueError):
mx.linalg.eigvalsh(mx.array([1.0, 2.0])) # 1D array
with self.assertRaises(ValueError):
mx.linalg.eigvalsh(
mx.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
) # Non-square matrix
if __name__ == "__main__":
unittest.main()