CPU mx.linalg.cholesky_inverse and mx.linalg.tri_inv (#1307)

* add cholesky inv + tri inv

* always run tri_inv on cpu

* consistent naming

python/src/linalg.cpp

@@ -261,6 +261,31 @@ void init_linalg(nb::module_& parent_module) {
             array: ``ainv`` such that ``dot(a, ainv) = dot(ainv, a) = eye(a.shape[0])``
       )pbdoc");
   m.def(
+      "tri_inv",
+      &tri_inv,
+      "a"_a,
+      "upper"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def tri_inv(a: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Compute the inverse of a triangular square matrix.
+
+        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.
+            upper (array): Whether the array is upper or lower triangular. Defaults to ``False``.
+            stream (Stream, optional): Stream or device. Defaults to ``None``
+              in which case the default stream of the default device is used.
+
+        Returns:
+            array: ``ainv`` such that ``dot(a, ainv) = dot(ainv, a) = eye(a.shape[0])``
+      )pbdoc");
+  m.def(
       "cholesky",
       &cholesky,
       "a"_a,
@@ -286,8 +311,46 @@ void init_linalg(nb::module_& parent_module) {
               in which case the default stream of the default device is used.
 
         Returns:
-          array: If ``upper = False``, it returns a lower trinagular ``L`` matrix such
+          array: If ``upper = False``, it returns a lower triangular ``L`` matrix such
           that ``dot(L, L.T) = a``.  If ``upper = True``, it returns an upper triangular
           ``U`` matrix such that ``dot(U.T, U) = a``.
       )pbdoc");
+  m.def(
+      "cholesky_inv",
+      &cholesky_inv,
+      "a"_a,
+      "upper"_a = false,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def cholesky_inv(L: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Compute the inverse of a real symmetric positive semi-definite matrix using it's Cholesky decomposition L.
+
+        Let A be a real symmetric positive semi-definite matrix and L its Cholesky definition such that:
+
+        .. math::
+
+          \begin{aligned}
+            \mathbf{A} = \mathbf{L}\mathbf{L}^T
+          \end{aligned}
+
+        This function computes :math:`\mathbf{A}^{-1}`.
+
+        This function supports arrays with at least 2 dimensions. When the input
+        has more than two dimensions, the Cholesky inverse is computed for each matrix
+        in the last two dimensions of ``L``.
+
+        If the input matrix is not a triangular matrix behaviour is undefined.
+
+        Args:
+            L (array): Input array.
+            upper (bool, optional): If ``True``, return the upper triangular Cholesky factor.
+              If ``False``, return the lower triangular Cholesky factor. Default: ``False``.
+            stream (Stream, optional): Stream or device. Defaults to ``None``
+              in which case the default stream of the default device is used.
+
+        Returns:
+          array: :math:`A^{-1}` where :math:`\mathbf{A} = \mathbf{L}\mathbf{L}^T`.
+      )pbdoc");
 }

python/tests/test_linalg.py

@@ -150,6 +150,20 @@ class TestLinalg(mlx_tests.MLXTestCase):
                 mx.allclose(M @ M_inv, mx.eye(M.shape[0]), rtol=0, atol=1e-5)
             )
 
+    def test_tri_inverse(self):
+        for upper in (False, True):
+            A = mx.array([[1, 0, 0], [6, -5, 0], [-9, 8, 7]], dtype=mx.float32)
+            B = mx.array([[7, 0, 0], [3, -2, 0], [1, 8, 3]], dtype=mx.float32)
+            if upper:
+                A = A.T
+                B = B.T
+            AB = mx.stack([A, B])
+            invs = mx.linalg.tri_inv(AB, upper=upper, stream=mx.cpu)
+            for M, M_inv in zip(AB, invs):
+                self.assertTrue(
+                    mx.allclose(M @ M_inv, mx.eye(M.shape[0]), rtol=0, atol=1e-5)
+                )
+
     def test_cholesky(self):
         sqrtA = mx.array(
             [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=mx.float32
@@ -167,6 +181,33 @@ 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_cholesky_inv(self):
+        mx.random.seed(7)
+
+        sqrtA = mx.array(
+            [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], dtype=mx.float32
+        )
+        A = sqrtA.T @ sqrtA / 81
+
+        N = 3
+        A = mx.random.uniform(shape=(N, N))
+        A = A @ A.T
+
+        for upper in (False, True):
+            L = mx.linalg.cholesky(A, upper=upper, stream=mx.cpu)
+            A_inv = mx.linalg.cholesky_inv(L, upper=upper, stream=mx.cpu)
+            self.assertTrue(mx.allclose(A @ A_inv, mx.eye(N), atol=1e-4))
+
+        # Multiple matrices
+        B = A + 1 / 9
+        AB = mx.stack([A, B])
+        Ls = mx.linalg.cholesky(AB, stream=mx.cpu)
+        for upper in (False, True):
+            Ls = mx.linalg.cholesky(AB, upper=upper, stream=mx.cpu)
+            AB_inv = mx.linalg.cholesky_inv(Ls, upper=upper, stream=mx.cpu)
+            for M, M_inv in zip(AB, AB_inv):
+                self.assertTrue(mx.allclose(M @ M_inv, mx.eye(N), atol=1e-4))
+
 
 if __name__ == "__main__":
     unittest.main()