zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Add the option for complex eigh

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-05-16 00:05:06 -07:00
parent 602f43e3d1
commit 9eb53248ba
5 changed files with 186 additions and 51 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