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

Add complex eigh (#2191)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-05-18 00:18:43 -07:00
committed by GitHub
parent 48ef3e74e2
commit 0654543dcc
5 changed files with 190 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/tests/test_linalg.py
View File

@@ -423,6 +423,13 @@ class TestLinalg(mlx_tests.MLXTestCase):
A_np = (A_np + np.transpose(A_np, (0, 2, 1))) / 2
check_eigs_and_vecs(A_np)
# Test with complex inputs
A_np = (
np.random.randn(8, 8, 2).astype(np.float32).view(np.complex64).squeeze(-1)
)
A_np = A_np + A_np.T.conj()
check_eigs_and_vecs(A_np)
# Test error cases
with self.assertRaises(ValueError):
mx.linalg.eigh(mx.array([1.0, 2.0])) # 1D array