Adding linalg.det

Author: Lucas Fernandes Martins

mlx/backend/cpu/luf.cpp

@@ -16,7 +16,8 @@ void luf_impl(
     array& lu,
     array& pivots,
     array& row_indices,
-    Stream stream) {
+    Stream stream,
+    bool allow_singular /* = false */) {
   int M = a.shape(-2);
   int N = a.shape(-1);
   int K = std::min(M, N);
@@ -54,50 +55,60 @@ void luf_impl(
   encoder.set_output_array(pivots);
   encoder.set_output_array(row_indices);
 
-  encoder.dispatch(
-      [a_ptr, pivots_ptr, row_indices_ptr, num_matrices, M, N, K]() mutable {
-        int info;
-        for (size_t i = 0; i < num_matrices; ++i) {
-          // Compute LU factorization of A
-          getrf<T>(
-              /* m */ &M,
-              /* n */ &N,
-              /* a */ a_ptr,
-              /* lda */ &M,
-              /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
-              /* info */ &info);
+  encoder.dispatch([a_ptr,
+                    pivots_ptr,
+                    row_indices_ptr,
+                    num_matrices,
+                    M,
+                    N,
+                    K,
+                    allow_singular]() mutable {
+    int info;
+    for (size_t i = 0; i < num_matrices; ++i) {
+      // Compute LU factorization of A
+      getrf<T>(
+          /* m */ &M,
+          /* n */ &N,
+          /* a */ a_ptr,
+          /* lda */ &M,
+          /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
+          /* info */ &info);
 
-          if (info != 0) {
-            std::stringstream ss;
-            ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info
-               << ((info > 0) ? " because matrix is singular"
-                              : " because argument had an illegal value");
-            throw std::runtime_error(ss.str());
-          }
+      if (info < 0) {
+        std::stringstream ss;
+        ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info
+           << " because argument had an illegal value";
+        throw std::runtime_error(ss.str());
+      } else if (info > 0 && !allow_singular) {
+        std::stringstream ss;
+        ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info
+           << " because matrix is singular";
+        throw std::runtime_error(ss.str());
+      }
 
-          // Subtract 1 to get 0-based index
-          int j = 0;
-          for (; j < K; ++j) {
-            pivots_ptr[j]--;
-            row_indices_ptr[j] = j;
-          }
-          for (; j < M; ++j) {
-            row_indices_ptr[j] = j;
-          }
-          for (int j = K - 1; j >= 0; --j) {
-            auto piv = pivots_ptr[j];
-            auto t1 = row_indices_ptr[piv];
-            auto t2 = row_indices_ptr[j];
-            row_indices_ptr[j] = t1;
-            row_indices_ptr[piv] = t2;
-          }
+      // Subtract 1 to get 0-based index
+      int j = 0;
+      for (; j < K; ++j) {
+        pivots_ptr[j]--;
+        row_indices_ptr[j] = j;
+      }
+      for (; j < M; ++j) {
+        row_indices_ptr[j] = j;
+      }
+      for (int j = K - 1; j >= 0; --j) {
+        auto piv = pivots_ptr[j];
+        auto t1 = row_indices_ptr[piv];
+        auto t2 = row_indices_ptr[j];
+        row_indices_ptr[j] = t1;
+        row_indices_ptr[piv] = t2;
+      }
 
-          // Advance pointers to the next matrix
-          a_ptr += M * N;
-          pivots_ptr += K;
-          row_indices_ptr += M;
-        }
-      });
+      // Advance pointers to the next matrix
+      a_ptr += M * N;
+      pivots_ptr += K;
+      row_indices_ptr += M;
+    }
+  });
 }
 
 void LUF::eval_cpu(
@@ -106,10 +117,22 @@ void LUF::eval_cpu(
   assert(inputs.size() == 1);
   switch (inputs[0].dtype()) {
     case float32:
-      luf_impl<float>(inputs[0], outputs[0], outputs[1], outputs[2], stream());
+      luf_impl<float>(
+          inputs[0],
+          outputs[0],
+          outputs[1],
+          outputs[2],
+          stream(),
+          allow_singular_);
       break;
     case float64:
-      luf_impl<double>(inputs[0], outputs[0], outputs[1], outputs[2], stream());
+      luf_impl<double>(
+          inputs[0],
+          outputs[0],
+          outputs[1],
+          outputs[2],
+          stream(),
+          allow_singular_);
       break;
     default:
       throw std::runtime_error(

mlx/linalg.cpp

@@ -583,7 +583,10 @@ void validate_lu(
   }
 }
 
-std::vector<array> lu_helper(const array& a, StreamOrDevice s /* = {} */) {
+std::vector<array> lu_helper(
+    const array& a,
+    StreamOrDevice s /* = {} */,
+    bool allow_singular = false) {
   int m = a.shape()[a.shape().size() - 2];
   int n = a.shape()[a.shape().size() - 1];
 
@@ -595,14 +598,14 @@ std::vector<array> lu_helper(const array& a, StreamOrDevice s /* = {} */) {
   return array::make_arrays(
       {a.shape(), pivots_shape, row_idx_shape},
       {a.dtype(), uint32, uint32},
-      std::make_shared<LUF>(to_stream(s)),
+      std::make_shared<LUF>(to_stream(s), allow_singular),
       {astype(a, a.dtype(), s)});
 }
 
 std::vector<array> lu(const array& a, StreamOrDevice s /* = {} */) {
   validate_lu(a, s, "[linalg::lu]");
 
-  auto out = lu_helper(a, s);
+  auto out = lu_helper(a, s, /*allow_singular=*/false);
   auto& LU = out[0];
   auto& row_pivots = out[2];
   auto L = tril(LU, /* k = */ -1, s);
@@ -705,4 +708,36 @@ array solve_triangular(
   return matmul(a_inv, b, s);
 }
 
+void validate_det(
+    const array& a,
+    const StreamOrDevice& stream,
+    const std::string& fname) {
+  check_cpu_stream(stream, fname);
+  check_float(a.dtype(), fname);
+  if (a.ndim() != 2 || a.shape(-2) != a.shape(-1)) {
+    std::ostringstream msg;
+    msg << fname
+        << " For determinant to be calculated, array must have 2 dimensions. Received array "
+           "with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
+array det(const array& a, StreamOrDevice s /* = {} */) {
+  validate_det(a, s, "[linalg::det]");
+
+  auto out = lu_helper(a, s, /*allow_singular=*/true);
+  auto& LU = out[0];
+  auto& pivots = out[1];
+
+  auto det_val = prod(diag(LU, 0, s), s);
+
+  auto indices = arange(pivots.shape(-1), pivots.dtype(), s);
+  auto num_swaps = sum(not_equal(pivots, indices, s), -1, false, s);
+  auto sign = power(array(-1.0f), astype(num_swaps, float32, s), s);
+
+  return multiply(sign, det_val, s);
+}
+
 } // namespace mlx::core::linalg

mlx/linalg.h

@@ -112,4 +112,6 @@ eigvalsh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
 MLX_API std::pair<array, array>
 eigh(const array& a, std::string UPLO = "L", StreamOrDevice s = {});
 
-} // namespace mlx::core::linalg
+MLX_API array det(const array& a, StreamOrDevice s = {});
+
+} // namespace mlx::core::linalg

mlx/primitives.h

@@ -2537,13 +2537,17 @@ class Eigh : public Primitive {
 /* LU Factorization primitive. */
 class LUF : public Primitive {
  public:
-  explicit LUF(Stream stream) : Primitive(stream) {}
+  explicit LUF(Stream stream, bool allow_singular = false)
+      : Primitive(stream), allow_singular_(allow_singular) {}
   void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
       override;
   void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
       override;
 
   DEFINE_NAME(LUF)
+
+ private:
+  bool allow_singular_{false};
 };
 
 } // namespace mlx::core

python/src/linalg.cpp

@@ -660,4 +660,31 @@ void init_linalg(nb::module_& parent_module) {
         Returns:
             array: The unique solution to the system ``AX = B``.
       )pbdoc");
+  m.def(
+      "det",
+      &mx::linalg::det,
+      "a"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def det(a: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Compute the determinant of a square matrix.
+
+        This function computes the determinant using LU factorization.
+        Singular matrices return a determinant of ``0``.
+
+        Args:
+            a (array): Input square matrix. Must be 2-D.
+            stream (Stream, optional): Stream or device. Defaults to ``None``
+              in which case the default stream of the default device is used.
+
+        Returns:
+            array: The determinant of the input matrix.
+
+        Example:
+            >>> a = mx.array([[1., 2.], [3., 4.]])
+            >>> mx.linalg.det(a)
+            array(-2, dtype=float32)
+      )pbdoc");
 }

python/tests/test_linalg.py

@@ -616,6 +616,85 @@ def test_solve_triangular(self):
         expected = np.linalg.solve(a, b)
         self.assertTrue(np.allclose(result, expected))
 
+    def test_det(self):
+        # 2x2 basic
+        a = mx.array([[1.0, 2.0], [3.0, 4.0]])
+        self.assertTrue(abs(mx.linalg.det(a).item() - (-2.0)) < 1e-4)
+
+        # 3x3 identity
+        a = mx.eye(3)
+        self.assertTrue(abs(mx.linalg.det(a).item() - 1.0) < 1e-4)
+
+        # 1x1
+        a = mx.array([[5.0]])
+        self.assertTrue(abs(mx.linalg.det(a).item() - 5.0) < 1e-4)
+
+        # 2x2 singular
+        a = mx.array([[1.0, 2.0], [2.0, 4.0]])
+        self.assertTrue(abs(mx.linalg.det(a).item()) < 1e-3)
+
+        # 3x3 general
+        a = mx.array([[1.0, 2.0, 3.0],
+                       [4.0, 5.0, 6.0],
+                       [7.0, 8.0, 0.0]])
+        self.assertTrue(abs(mx.linalg.det(a).item() - 27.0) < 1e-4)
+
+        # 3x3 diagonal
+        a = mx.array([[2.0, 0.0, 0.0],
+                       [0.0, 3.0, 0.0],
+                       [0.0, 0.0, 4.0]])
+        self.assertTrue(abs(mx.linalg.det(a).item() - 24.0) < 1e-4)
+
+        # 3x3 upper triangular
+        a = mx.array([[2.0, 1.0, 3.0],
+                       [0.0, 5.0, 7.0],
+                       [0.0, 0.0, 4.0]])
+        self.assertTrue(abs(mx.linalg.det(a).item() - 40.0) < 1e-4)
+
+        # 2x2 permutation (negative det)
+        a = mx.array([[0.0, 1.0], [1.0, 0.0]])
+        self.assertTrue(abs(mx.linalg.det(a).item() - (-1.0)) < 1e-4)
+
+        # 4x4 general
+        a = mx.array([[1.0, 0.0, 2.0, -1.0],
+                       [3.0, 0.0, 0.0,  5.0],
+                       [2.0, 1.0, 4.0, -3.0],
+                       [1.0, 0.0, 5.0,  0.0]])
+        self.assertTrue(abs(mx.linalg.det(a).item() - 30.0) < 1e-4)
+
+        # 4x4 scaled identity
+        a = 3.0 * mx.eye(4)
+        self.assertTrue(abs(mx.linalg.det(a).item() - 81.0) < 1e-4)
+
+        # 3x3 zeros (singular)
+        a = mx.zeros((3, 3))
+        self.assertTrue(abs(mx.linalg.det(a).item()) < 1e-3)
+
+        # 2x2 all negative
+        a = mx.array([[-1.0, -2.0], [-3.0, -4.0]])
+        self.assertTrue(abs(mx.linalg.det(a).item() - (-2.0)) < 1e-4)
+
+        # float64
+        a = mx.array([[1.0, 2.0], [3.0, 4.0]], dtype=mx.float64)
+        self.assertTrue(abs(mx.linalg.det(a).item() - (-2.0)) < 1e-10)
+
+    def test_det_throws(self):
+        # 1D array
+        with self.assertRaises(Exception):
+            mx.eval(mx.linalg.det(mx.array([1.0, 2.0, 3.0])))
+
+        # Non-square matrix
+        with self.assertRaises(Exception):
+            mx.eval(mx.linalg.det(mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])))
+
+        # 0D scalar
+        with self.assertRaises(Exception):
+            mx.eval(mx.linalg.det(mx.array(5.0)))
+
+        # Integer input
+        with self.assertRaises(Exception):
+            mx.eval(mx.linalg.det(mx.array([[1, 2], [3, 4]])))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()