Pinv (#875)

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");
 }

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)