From e6b223df5f2f02d85f7fd4a144a204c6dee437d1 Mon Sep 17 00:00:00 2001 From: Aditya Dhulipala Date: Tue, 27 Aug 2024 23:06:12 -0700 Subject: [PATCH] Pinv (#875) --- mlx/linalg.cpp | 43 ++++++++++++++++++++++++++++++++++++ mlx/linalg.h | 2 ++ python/src/linalg.cpp | 24 ++++++++++++++++++++ python/tests/test_linalg.py | 12 ++++++++++ tests/linalg_tests.cpp | 44 ++++++++++++++++++++++++++++++++++++- 5 files changed, 124 insertions(+), 1 deletion(-) diff --git a/mlx/linalg.cpp b/mlx/linalg.cpp index f59efd5ad..9a397b868 100644 --- a/mlx/linalg.cpp +++ b/mlx/linalg.cpp @@ -306,6 +306,49 @@ array cholesky( {a}); } +array pinv(const array& a, StreamOrDevice s /* = {} */) { + if (a.dtype() != float32) { + std::ostringstream msg; + msg << "[linalg::pinv] Arrays must type float32. Received array " + << "with type " << a.dtype() << "."; + throw std::invalid_argument(msg.str()); + } + if (a.ndim() < 2) { + std::ostringstream msg; + msg << "[linalg::pinv] Arrays must have >= 2 dimensions. Received array " + << "with " << a.ndim() << " dimensions."; + throw std::invalid_argument(msg.str()); + } + + int m = a.shape(-2); + int n = a.shape(-1); + int k = std::min(m, n); + auto outs = linalg::svd(a, s); + array U = outs[0]; + array S = outs[1]; + array V = outs[2]; + + std::vector starts(a.ndim(), 0); + std::vector ends = a.shape(); + int i = a.ndim() - 2; + int j = a.ndim() - 1; + + // Prepare U + ends[i] = m; + ends[j] = k; + U = swapaxes(slice(U, starts, ends, s), -1, -2, s); + + // Prepare V + ends[i] = k; + ends[j] = n; + V = swapaxes(slice(V, starts, ends, s), -1, -2, s); + + // Prepare S + S = expand_dims(S, -2, s); + + return matmul(divide(V, S, s), U); +} + array cholesky_inv( const array& L, bool upper /* = false */, diff --git a/mlx/linalg.h b/mlx/linalg.h index 6f3eb33fc..3ffca476c 100644 --- a/mlx/linalg.h +++ b/mlx/linalg.h @@ -70,6 +70,8 @@ array tri_inv(const array& a, bool upper = false, StreamOrDevice s = {}); array cholesky(const array& a, bool upper = false, StreamOrDevice s = {}); +array pinv(const array& a, StreamOrDevice s = {}); + array cholesky_inv(const array& a, bool upper = false, StreamOrDevice s = {}); } // namespace mlx::core::linalg diff --git a/python/src/linalg.cpp b/python/src/linalg.cpp index b17ba185e..235e4f828 100644 --- a/python/src/linalg.cpp +++ b/python/src/linalg.cpp @@ -353,4 +353,28 @@ void init_linalg(nb::module_& parent_module) { Returns: array: :math:`\mathbf{A^{-1}}` where :math:`\mathbf{A} = \mathbf{L}\mathbf{L}^T`. )pbdoc"); + m.def( + "pinv", + &pinv, + "a"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def pinv(a: array, *, stream: Union[None, Stream, Device] = None) -> array"), + R"pbdoc( + Compute the (Moore-Penrose) pseudo-inverse of a matrix. + + This function calculates a generalized inverse of a matrix using its + singular-value decomposition. This function supports arrays with at least 2 dimensions. + When the input has more than two dimensions, the inverse is computed for each + matrix in the last two dimensions of ``a``. + + Args: + a (array): 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: ``aplus`` such that ``a @ aplus @ a = a`` + )pbdoc"); } diff --git a/python/tests/test_linalg.py b/python/tests/test_linalg.py index 0a6fe9a53..294051077 100644 --- a/python/tests/test_linalg.py +++ b/python/tests/test_linalg.py @@ -181,6 +181,18 @@ class TestLinalg(mlx_tests.MLXTestCase): for M, L in zip(AB, Ls): self.assertTrue(mx.allclose(L @ L.T, M, rtol=1e-5, atol=1e-7)) + def test_pseudo_inverse(self): + A = mx.array([[1, 2, 3], [6, -5, 4], [-9, 8, 7]], dtype=mx.float32) + A_plus = mx.linalg.pinv(A, stream=mx.cpu) + self.assertTrue(mx.allclose(A @ A_plus @ A, A, rtol=0, atol=1e-5)) + + # Multiple matrices + B = A - 100 + AB = mx.stack([A, B]) + pinvs = mx.linalg.pinv(AB, stream=mx.cpu) + for M, M_plus in zip(AB, pinvs): + self.assertTrue(mx.allclose(M @ M_plus @ M, M, rtol=0, atol=1e-3)) + def test_cholesky_inv(self): mx.random.seed(7) diff --git a/tests/linalg_tests.cpp b/tests/linalg_tests.cpp index 2af868965..186bbc613 100644 --- a/tests/linalg_tests.cpp +++ b/tests/linalg_tests.cpp @@ -347,4 +347,46 @@ TEST_CASE("test matrix cholesky") { .item()); CHECK(allclose(matmul(transpose(U), U), A, /* rtol = */ 0, /* atol = */ 1e-6) .item()); -} \ No newline at end of file +} + +TEST_CASE("test matrix pseudo-inverse") { + // 0D and 1D throw + CHECK_THROWS(linalg::pinv(array(0.0), Device::cpu)); + CHECK_THROWS(linalg::pinv(array({0.0, 1.0}), Device::cpu)); + + // Unsupported types throw + CHECK_THROWS(linalg::pinv(array({0, 1}, {1, 2}), Device::cpu)); + + { // Square m == n + const auto A = array({1.0, 2.0, 3.0, 4.0}, {2, 2}); + const auto A_pinv = linalg::pinv(A, Device::cpu); + const auto A_again = matmul(matmul(A, A_pinv), A); + CHECK(allclose(A_again, A).item()); + const auto A_pinv_again = matmul(matmul(A_pinv, A), A_pinv); + CHECK(allclose(A_pinv_again, A_pinv).item()); + } + { // Rectangular matrix m < n + const auto prng_key = random::key(42); + const auto A = random::normal({4, 5}, prng_key); + const auto A_pinv = linalg::pinv(A, Device::cpu); + const auto zeros = zeros_like(A_pinv, Device::cpu); + CHECK_FALSE(allclose(zeros, A_pinv, /* rtol = */ 0, /* atol = */ 1e-6) + .item()); + const auto A_again = matmul(matmul(A, A_pinv), A); + CHECK(allclose(A_again, A).item()); + const auto A_pinv_again = matmul(matmul(A_pinv, A), A_pinv); + CHECK(allclose(A_pinv_again, A_pinv).item()); + } + { // Rectangular matrix m > n + const auto prng_key = random::key(10); + const auto A = random::normal({6, 5}, prng_key); + const auto A_pinv = linalg::pinv(A, Device::cpu); + const auto zeros2 = zeros_like(A_pinv, Device::cpu); + CHECK_FALSE(allclose(zeros2, A_pinv, /* rtol = */ 0, /* atol = */ 1e-6) + .item()); + const auto A_again = matmul(matmul(A, A_pinv), A); + CHECK(allclose(A_again, A).item()); + const auto A_pinv_again = matmul(matmul(A_pinv, A), A_pinv); + CHECK(allclose(A_pinv_again, A_pinv).item()); + } +}