L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add determinant and sign-log-determinant functions to mlx.core.linalg (#3416)

Browse Source 
Co-authored-by: Lucas Fernandes Martins <[email protected]>
This commit is contained in:
Abhilash Shankarampeta
2026-05-05 09:06:23 +09:00
committed by GitHub
co-authored by Lucas Fernandes Martins
parent e8ebdebeeb
commit 0938db7e54
7 changed files with 568 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/linalg.cpp
+73
View File
@@ -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");
}
python/tests/test_linalg.py
+255
View File
@@ -520,6 +520,19 @@ class TestLinalg(mlx_tests.MLXTestCase):
P, L, U = mx.linalg.lu(a, stream=mx.cpu)
self.assertTrue(mx.allclose(L[P, :] @ U, a))
# Test singular matrix (should not throw)
a = mx.array(
[
[1.0, 2.0, 3.0, 4.0],
[2.0, 4.0, 6.0, 8.0],
[0.0, 1.0, 1.0, 0.0],
[1.0, 0.0, 0.0, 1.0],
]
)
P, L, U = mx.linalg.lu(a, stream=mx.cpu)
L_permuted = mx.take_along_axis(L, P[..., None], axis=-2)
self.assertTrue(mx.allclose(L_permuted @ U, a))
def test_lu_factor(self):
mx.random.seed(7)
@@ -616,6 +629,248 @@ class TestLinalg(mlx_tests.MLXTestCase):
expected = np.linalg.solve(a, b)
self.assertTrue(np.allclose(result, expected))
def test_det(self):
# 1x1 fast path
A = mx.array([[5.0]])
self.assertTrue(np.allclose(mx.linalg.det(A, stream=mx.cpu), 5.0))
# 2x2 fast path
A = mx.array([[1.0, 2.0], [3.0, 4.0]])
d = mx.linalg.det(A, stream=mx.cpu)
self.assertTrue(np.allclose(d, -2.0))
# 3x3 fast path
A = mx.array([[1.0, 2.0, 3.0], [0.0, 1.0, 4.0], [5.0, 6.0, 0.0]])
d = mx.linalg.det(A, stream=mx.cpu)
expected = np.linalg.det(np.array(A))
self.assertTrue(np.allclose(d, expected, atol=1e-5))
# 4x4 LU path: compare with numpy
np.random.seed(42)
A_np = np.random.randn(4, 4).astype(np.float32)
A_mx = mx.array(A_np)
d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
d_np = np.linalg.det(A_np)
self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
# 5x5 LU path
A_np = np.random.randn(5, 5).astype(np.float32)
A_mx = mx.array(A_np)
d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
d_np = np.linalg.det(A_np)
self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
# Identity matrix
A = mx.eye(5)
self.assertTrue(np.allclose(mx.linalg.det(A, stream=mx.cpu), 1.0))
# Batched: (3, 4, 4)
A_np = np.random.randn(3, 4, 4).astype(np.float32)
A_mx = mx.array(A_np)
d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
d_np = np.linalg.det(A_np)
self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
# Multi-batch: (2, 3, 3, 3)
A_np = np.random.randn(2, 3, 3, 3).astype(np.float32)
A_mx = mx.array(A_np)
d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
d_np = np.linalg.det(A_np)
self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4))
# Integer input auto-promotes to float
A = mx.array([[1, 2], [3, 4]])
d = mx.linalg.det(A, stream=mx.cpu)
self.assertTrue(np.allclose(d, -2.0))
# float64
A_np = np.random.randn(4, 4).astype(np.float64)
A_mx = mx.array(A_np)
d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
d_np = np.linalg.det(A_np)
self.assertTrue(np.allclose(d_mx, d_np, atol=1e-10))
# Singular 4x4 matrix (LU path): det should be 0
A = mx.array(
[
[1.0, 2.0, 3.0, 4.0],
[2.0, 4.0, 6.0, 8.0],
[0.0, 1.0, 1.0, 0.0],
[1.0, 0.0, 0.0, 1.0],
]
)
d = mx.linalg.det(A, stream=mx.cpu)
self.assertTrue(np.allclose(d, 0.0, atol=1e-5))
# Singular 5x5 matrix (LU path)
A_np = np.ones((5, 5), dtype=np.float32)
A_mx = mx.array(A_np)
d = mx.linalg.det(A_mx, stream=mx.cpu)
self.assertTrue(np.allclose(d, 0.0, atol=1e-5))
# Batched singular matrices (LU path)
A_np = np.array([np.diag([1.0, 2.0, 0.0, 3.0]), np.eye(4, dtype=np.float32)])
A_mx = mx.array(A_np)
d_mx = mx.linalg.det(A_mx, stream=mx.cpu)
d_np = np.linalg.det(A_np)
self.assertTrue(np.allclose(d_mx, d_np, atol=1e-5))
# Empty 0x0 matrix: det is the empty product = 1
d = mx.linalg.det(mx.zeros((0, 0)), stream=mx.cpu)
self.assertEqual(d.shape, ())
self.assertEqual(float(d), 1.0)
# Batched empty matrices: shape preserves batch dims
d = mx.linalg.det(mx.zeros((3, 0, 0)), stream=mx.cpu)
self.assertTrue(np.allclose(d, np.linalg.det(np.zeros((3, 0, 0)))))
# Error: non-square
with self.assertRaises(ValueError):
mx.linalg.det(mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), stream=mx.cpu)
# Error: 1D
with self.assertRaises(ValueError):
mx.linalg.det(mx.array([1.0, 2.0]), stream=mx.cpu)
# Error: complex unsupported (small-matrix path)
with self.assertRaises(ValueError):
mx.linalg.det(mx.array([[1.0 + 1j, 2.0], [3.0, 4.0]]), stream=mx.cpu)
# Error: complex unsupported (LU path)
with self.assertRaises(ValueError):
mx.linalg.det(mx.eye(4).astype(mx.complex64), stream=mx.cpu)
def test_slogdet(self):
# 2x2: det = -2 => sign = -1, logabsdet = log(2)
A = mx.array([[1.0, 2.0], [3.0, 4.0]])
sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
self.assertTrue(np.allclose(sign, -1.0))
self.assertTrue(np.allclose(logabsdet, np.log(2.0), atol=1e-5))
# Identity: sign = 1, logabsdet = 0
A = mx.eye(4)
sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
self.assertTrue(np.allclose(sign, 1.0))
self.assertTrue(np.allclose(logabsdet, 0.0, atol=1e-6))
# Compare with numpy for random matrices
np.random.seed(42)
for n in [1, 2, 3, 4, 5]:
A_np = np.random.randn(n, n).astype(np.float32)
A_mx = mx.array(A_np)
sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
sign_np, logabs_np = np.linalg.slogdet(A_np)
with self.subTest(n=n):
self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4))
# Singular matrix 2x2 (fast path): sign = 0, logabsdet = -inf
A = mx.array([[1.0, 2.0], [2.0, 4.0]])
sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
self.assertEqual(float(sign), 0.0)
self.assertEqual(float(logabsdet), float("-inf"))
# Singular 4x4 matrix (LU path): sign = 0, logabsdet = -inf
A = mx.array(
[
[1.0, 2.0, 3.0, 4.0],
[2.0, 4.0, 6.0, 8.0],
[0.0, 1.0, 1.0, 0.0],
[1.0, 0.0, 0.0, 1.0],
]
)
sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
self.assertEqual(float(sign), 0.0)
self.assertEqual(float(logabsdet), float("-inf"))
# Singular 5x5 matrix (LU path): all-ones matrix
A = mx.array(np.ones((5, 5), dtype=np.float32))
sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
self.assertEqual(float(sign), 0.0)
self.assertEqual(float(logabsdet), float("-inf"))
# Batched with mix of singular and non-singular (LU path)
A_np = np.array([np.diag([1.0, 2.0, 0.0, 3.0]), np.eye(4, dtype=np.float32)])
A_mx = mx.array(A_np)
sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
sign_np, logabs_np = np.linalg.slogdet(A_np)
self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
# Check -inf for singular, 0.0 for identity
self.assertEqual(float(logabs_mx[0]), float("-inf"))
self.assertTrue(np.allclose(logabs_mx[1], 0.0, atol=1e-6))
# Batched
A_np = np.random.randn(3, 4, 4).astype(np.float32)
A_mx = mx.array(A_np)
sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
sign_np, logabs_np = np.linalg.slogdet(A_np)
self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4))
# Multi-batch
A_np = np.random.randn(2, 3, 3, 3).astype(np.float32)
A_mx = mx.array(A_np)
sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
sign_np, logabs_np = np.linalg.slogdet(A_np)
self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5))
self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4))
# Numerical stability: large matrix where det overflows
# 0.1 * I_100 has det = 0.1^100 which underflows in float32
# but slogdet should give sign=1, logabsdet = 100*log(0.1)
n = 100
A = mx.array(0.1) * mx.eye(n)
sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu)
self.assertTrue(np.allclose(sign, 1.0))
self.assertTrue(np.allclose(logabsdet, n * np.log(0.1), atol=1e-3))
# Verify det = sign * exp(logabsdet) for non-singular cases
A_np = np.random.randn(5, 5).astype(np.float32)
A_mx = mx.array(A_np)
sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
det_mx = mx.linalg.det(A_mx, stream=mx.cpu)
reconstructed = float(sign_mx) * np.exp(float(logabs_mx))
self.assertTrue(np.allclose(float(det_mx), reconstructed, rtol=1e-4))
# float64
A_np = np.random.randn(4, 4).astype(np.float64)
A_mx = mx.array(A_np)
sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu)
sign_np, logabs_np = np.linalg.slogdet(A_np)
self.assertTrue(np.allclose(sign_mx, sign_np))
self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-10))
# Empty 0x0 matrix: sign = 1, logabsdet = 0 (empty product)
sign, logabsdet = mx.linalg.slogdet(mx.zeros((0, 0)), stream=mx.cpu)
self.assertEqual(sign.shape, ())
self.assertEqual(logabsdet.shape, ())
self.assertEqual(float(sign), 1.0)
self.assertEqual(float(logabsdet), 0.0)
# Batched empty matrices
sign, logabsdet = mx.linalg.slogdet(mx.zeros((3, 0, 0)), stream=mx.cpu)
sign_np, logabs_np = np.linalg.slogdet(np.zeros((3, 0, 0)))
self.assertTrue(np.allclose(sign, sign_np))
self.assertTrue(np.allclose(logabsdet, logabs_np))
# Error: non-square
with self.assertRaises(ValueError):
mx.linalg.slogdet(
mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), stream=mx.cpu
)
# Error: 1D
with self.assertRaises(ValueError):
mx.linalg.slogdet(mx.array([1.0, 2.0]), stream=mx.cpu)
# Error: complex unsupported (small-matrix path)
with self.assertRaises(ValueError):
mx.linalg.slogdet(mx.array([[1.0 + 1j, 2.0], [3.0, 4.0]]), stream=mx.cpu)
# Error: complex unsupported (LU path)
with self.assertRaises(ValueError):
mx.linalg.slogdet(mx.eye(4).astype(mx.complex64), stream=mx.cpu)
if __name__ == "__main__":
mlx_tests.MLXTestRunner()