diff --git a/mlx/backend/cpu/eigh.cpp b/mlx/backend/cpu/eigh.cpp index d457c1fd99..88bcc87e2e 100644 --- a/mlx/backend/cpu/eigh.cpp +++ b/mlx/backend/cpu/eigh.cpp @@ -155,10 +155,12 @@ void eigh_impl( auto& encoder = cpu::get_command_encoder(stream); encoder.set_output_array(vectors); encoder.set_output_array(values); + // LAPACK reads the row-major input as its (conjugate) transpose, so the + // requested triangle is the opposite one in LAPACK's view. encoder.dispatch([vec_ptr, eig_ptr, jobz, - uplo = uplo[0], + uplo = uplo[0] == 'L' ? 'U' : 'L', N = vectors.shape(-1), size = vectors.size()]() mutable { // Work query @@ -202,6 +204,15 @@ void Eigh::eval_cpu( a.flags().row_contiguous ? CopyType::Vector : CopyType::General, stream()); + // Nothing to decompose for n = 0; LAPACK rejects lda = 0. + if (a.shape(-1) == 0) { + if (!compute_eigenvectors_) { + auto& encoder = cpu::get_command_encoder(stream()); + encoder.add_temporary(vectors); + } + return; + } + if (compute_eigenvectors_) { // Set the strides and flags so the eigenvectors // are in the columns of the output diff --git a/mlx/backend/cpu/svd.cpp b/mlx/backend/cpu/svd.cpp index ca01a0a65a..808ff92787 100644 --- a/mlx/backend/cpu/svd.cpp +++ b/mlx/backend/cpu/svd.cpp @@ -206,6 +206,14 @@ void svd_impl( using R = typename SVDWork::R; + // Nothing to decompose for empty matrices; LAPACK rejects lda = 0. + if (M == 0 || N == 0) { + for (auto& o : outputs) { + o.set_data(allocator::malloc(o.nbytes())); + } + return; + } + size_t num_matrices = a.size() / (M * N); // lapack clobbers the input, so we have to make a copy. diff --git a/python/tests/test_linalg.py b/python/tests/test_linalg.py index afdf75d7c8..fc66dd2b14 100644 --- a/python/tests/test_linalg.py +++ b/python/tests/test_linalg.py @@ -188,6 +188,13 @@ def test_svd_decomposition(self): ) ) + # Zero-size inputs + for shape in [(0, 4, 4), (3, 0, 0)]: + U, S, Vt = mx.linalg.svd(mx.zeros(shape), stream=mx.cpu) + mx.eval(U, S, Vt) + self.assertEqual(S.shape, shape[:-1]) + self.assertEqual(U.shape, shape) + # Test float64 - use CPU stream since float64 is not supported on GPU with mx.stream(mx.cpu): A_f64 = mx.array( @@ -493,6 +500,21 @@ def check_eigs_and_vecs(A_np, kwargs={}): A_np = A_np + A_np.T.conj() check_eigs_and_vecs(A_np) + # UPLO picks the triangle like numpy; only observable when the two + # triangles disagree + A_np = np.array([[1.0, 999.0], [2.0, 3.0]], dtype=np.float32) + for uplo in ("L", "U"): + w = mx.linalg.eigvalsh(mx.array(A_np), UPLO=uplo, stream=mx.cpu) + w_np = np.linalg.eigvalsh(A_np, UPLO=uplo) + self.assertTrue(np.allclose(w, w_np, atol=1e-5)) + + # Zero-size inputs + for shape in [(0, 4, 4), (3, 0, 0)]: + w, v = mx.linalg.eigh(mx.zeros(shape), stream=mx.cpu) + mx.eval(w, v) + self.assertEqual(w.shape, shape[:-1]) + self.assertEqual(v.shape, shape) + # Test error cases with self.assertRaises(ValueError): mx.linalg.eigh(mx.array([1.0, 2.0])) # 1D array