Add matrix inversion primitive (#822)

python/src/linalg.cpp

@@ -241,4 +241,27 @@ void init_linalg(py::module_& parent_module) {
             tuple(array, array, array): The ``U``, ``S``, and ``Vt`` matrices, such that
               ``A = U @ diag(S) @ Vt``
       )pbdoc");
+  m.def(
+      "inv",
+      &inv,
+      "a"_a,
+      py::kw_only(),
+      "stream"_a = none,
+      R"pbdoc(
+        inv(a: array, *, stream: Union[None, Stream, Device] = None) -> array
+
+        Compute the inverse of a 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.
+            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");
 }

python/tests/test_linalg.py

@@ -136,6 +136,20 @@ class TestLinalg(mlx_tests.MLXTestCase):
                 mx.allclose(U[:, : len(S)] @ mx.diag(S) @ Vt, M, rtol=1e-5, atol=1e-7)
             )
 
+    def test_inverse(self):
+        A = mx.array([[1, 2, 3], [6, -5, 4], [-9, 8, 7]], dtype=mx.float32)
+        A_inv = mx.linalg.inv(A, stream=mx.cpu)
+        self.assertTrue(mx.allclose(A @ A_inv, mx.eye(A.shape[0]), rtol=0, atol=1e-6))
+
+        # Multiple matrices
+        B = A - 100
+        AB = mx.stack([A, B])
+        invs = mx.linalg.inv(AB, 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)
+            )
+
 
 if __name__ == "__main__":
     unittest.main()