Add determinant and sign-log-determinant functions to mlx.core.linalg

docs/src/python/linalg.rst

@@ -14,6 +14,7 @@ Linear Algebra
     cholesky
     cholesky_inv
     cross
+    det
     qr
     svd
     eigvals
@@ -23,5 +24,6 @@ Linear Algebra
     lu
     lu_factor
     pinv
+    slogdet
     solve
     solve_triangular

mlx/backend/cpu/luf.cpp

@@ -67,11 +67,10 @@ void luf_impl(
               /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
               /* info */ &info);
 
-          if (info != 0) {
+          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");
+               << " because argument had an illegal value";
             throw std::runtime_error(ss.str());
           }
 

mlx/linalg.cpp

@@ -705,4 +705,162 @@ array solve_triangular(
   return matmul(a_inv, b, s);
 }
 
-} // namespace mlx::core::linalg
+void validate_det(
+    const array& a,
+    const StreamOrDevice& stream,
+    const std::string& fname) {
+  check_cpu_stream(stream, fname);
+
+  if (a.ndim() < 2) {
+    std::ostringstream msg;
+    msg << fname
+        << " Arrays must have >= 2 dimensions. Received array "
+           "with "
+        << a.ndim() << " dimensions.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (a.shape(-1) != a.shape(-2)) {
+    throw std::invalid_argument(fname + " Only defined for square matrices.");
+  }
+}
+
+array det_raw_small(const array& a, StreamOrDevice s) {
+  int n = a.shape(-1);
+
+  // Helper to extract a[..., i, j] from the last two dims
+  auto elem = [&](int i, int j) {
+    auto starts = Shape(a.ndim(), 0);
+    auto stops = a.shape();
+    starts[a.ndim() - 2] = i;
+    stops[a.ndim() - 2] = i + 1;
+    starts[a.ndim() - 1] = j;
+    stops[a.ndim() - 1] = j + 1;
+    return squeeze(squeeze(slice(a, starts, stops, s), -1, s), -1, s);
+  };
+
+  if (n == 1) {
+    return elem(0, 0);
+  } else if (n == 2) {
+    return subtract(
+        multiply(elem(0, 0), elem(1, 1), s),
+        multiply(elem(0, 1), elem(1, 0), s),
+        s);
+  } else {
+    // 3x3: a00*(a11*a22 - a12*a21) - a01*(a10*a22 - a12*a20) + a02*(a10*a21 -
+    // a11*a20)
+    auto a00 = elem(0, 0), a01 = elem(0, 1), a02 = elem(0, 2);
+    auto a10 = elem(1, 0), a11 = elem(1, 1), a12 = elem(1, 2);
+    auto a20 = elem(2, 0), a21 = elem(2, 1), a22 = elem(2, 2);
+    return add(
+        subtract(
+            multiply(
+                a00,
+                subtract(multiply(a11, a22, s), multiply(a12, a21, s), s),
+                s),
+            multiply(
+                a01,
+                subtract(multiply(a10, a22, s), multiply(a12, a20, s), s),
+                s),
+            s),
+        multiply(
+            a02, subtract(multiply(a10, a21, s), multiply(a11, a20, s), s), s),
+        s);
+  }
+}
+
+std::pair<array, array> slogdet_impl(const array& input, StreamOrDevice s) {
+  int n = input.shape(-1);
+  auto dtype = input.dtype();
+
+  // Small-matrix fast path
+  if (n <= 3) {
+    auto raw = det_raw_small(input, s);
+    auto abs_raw = abs(raw, s);
+    auto sgn = sign(raw, s);
+    auto logabs = log(abs_raw, s);
+    return std::make_pair(sgn, logabs);
+  }
+
+  // General LU-based path
+  auto [LU, pivots] = lu_factor(input, s);
+
+  // Extract diagonal of U
+  auto diag = diagonal(LU, 0, -2, -1, s);
+
+  // Permutation parity: count positions where pivot[i] != i
+  int k = std::min(input.shape(-2), input.shape(-1));
+  auto iota = arange(0, k, uint32, s);
+  auto parity = astype(
+      sum(not_equal(pivots, iota, s),
+          /* axis = */ -1,
+          /* keepdims = */ false,
+          s),
+      int32,
+      s);
+
+  // Count negative diagonal elements
+  auto num_neg = astype(
+      sum(less(diag, array(0.0f, dtype), s),
+          /* axis = */ -1,
+          /* keepdims = */ false,
+          s),
+      int32,
+      s);
+
+  // sign = (-1)^(parity + num_neg)
+  auto total = add(parity, num_neg, s);
+  auto sign_val = astype(
+      subtract(
+          array(1, int32),
+          multiply(array(2, int32), remainder(total, array(2, int32), s), s),
+          s),
+      dtype,
+      s);
+
+  // logabsdet = sum(log(abs(diag)))
+  auto logabsdet =
+      sum(log(abs(diag, s), s), /* axis = */ -1, /* keepdims = */ false, s);
+
+  // Handle singular matrices: any zero on diagonal
+  auto is_zero =
+      any(equal(diag, array(0.0f, dtype), s),
+          /* axis = */ -1,
+          /* keepdims = */ false,
+          s);
+  sign_val = where(is_zero, array(0.0f, dtype), sign_val, s);
+  logabsdet = where(
+      is_zero,
+      array(-std::numeric_limits<float>::infinity(), dtype),
+      logabsdet,
+      s);
+
+  return std::make_pair(sign_val, logabsdet);
+}
+
+std::pair<array, array> slogdet(const array& a, StreamOrDevice s /* = {} */) {
+  validate_det(a, s, "[linalg::slogdet]");
+
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  return slogdet_impl(input, s);
+}
+
+array det(const array& a, StreamOrDevice s /* = {} */) {
+  validate_det(a, s, "[linalg::det]");
+
+  auto dtype = at_least_float(a.dtype());
+  auto input = astype(a, dtype, s);
+  int n = input.shape(-1);
+
+  // Small-matrix fast path: compute directly, skip log/exp round-trip
+  if (n <= 3) {
+    return det_raw_small(input, s);
+  }
+
+  // General case: det = sign * exp(logabsdet)
+  auto [sign_val, logabsdet] = slogdet_impl(input, s);
+  return multiply(sign_val, exp(logabsdet, s), s);
+}
+
+} // namespace mlx::core::linalg

mlx/linalg.h

@@ -112,4 +112,8 @@ 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 = {});
 
+MLX_API array det(const array& a, StreamOrDevice s = {});
+
+MLX_API std::pair<array, array> slogdet(const array& a, StreamOrDevice s = {});
+
 } // namespace mlx::core::linalg

python/src/linalg.cpp

@@ -660,4 +660,77 @@ 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 supports arrays with at least 2 dimensions. When the
+        input has more than two dimensions, the determinant is computed for
+        each matrix in the last two dimensions.
+
+        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: The determinant(s) of the input matrix (matrices).
+
+        Example:
+            >>> A = mx.array([[1., 2.], [3., 4.]])
+            >>> mx.linalg.det(A, stream=mx.cpu)
+            array(-2, dtype=float32)
+      )pbdoc");
+
+  m.def(
+      "slogdet",
+      [](const mx::array& a, mx::StreamOrDevice s) {
+        auto result = mx::linalg::slogdet(a, s);
+        return nb::make_tuple(result.first, result.second);
+      },
+      "a"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def slogdet(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"),
+      R"pbdoc(
+        Compute the sign and natural log of the absolute value of the
+        determinant of a square matrix.
+
+        This function supports arrays with at least 2 dimensions. When the
+        input has more than two dimensions, the sign and log-absolute-determinant
+        are computed for each matrix in the last two dimensions.
+
+        For a singular matrix, ``sign`` is 0 and ``logabsdet`` is ``-inf``.
+
+        The determinant can be reconstructed as ``det = sign * exp(logabsdet)``.
+        This is more numerically stable than computing the determinant directly
+        for matrices with large or small determinants.
+
+        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:
+            tuple(array, array): The ``sign`` and ``logabsdet`` of the
+              determinant. ``sign`` is -1, 0, or +1. ``logabsdet`` is the
+              natural log of the absolute value of the determinant.
+
+        Example:
+            >>> A = mx.array([[1., 2.], [3., 4.]])
+            >>> sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
+            >>> sign
+            array(-1, dtype=float32)
+            >>> logabsdet
+            array(0.693147, dtype=float32)
+      )pbdoc");
 }