zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 17:31:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Enforce triangular matrix form in tri_inv (#1876)

Browse Source 
* fix tri_inv bug

* Revert "fix tri_inv bug"

This reverts commit b74b290201.

* Make sure that tri_inv returns a triangular matrix

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
This commit is contained in:
Abe Leininger 2025-02-19 14:42:33 -06:00 committed by GitHub
parent 71de73a668
commit 344a29506e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 23 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
mlx/backend/cpu/inverse.cpp
View File

@ -81,7 +81,22 @@ void general_inv(array& inv, int N, int i) {
void tri_inv(array& inv, int N, int i, bool upper) {
const char uplo = upper ? 'L' : 'U';
const char diag = 'N';
int info = strtri_wrapper(uplo, diag, inv.data<float>() + N * N * i, N);
float* data = inv.data<float>() + N * N * i;
int info = strtri_wrapper(uplo, diag, data, N);
// zero out the other triangle
if (upper) {
for (int i = 0; i < N; i++) {
std::fill(data, data + i, 0.0f);
data += N;
}
} else {
for (int i = 0; i < N; i++) {
std::fill(data + i + 1, data + N, 0.0f);
data += N;
}
}
if (info != 0) {
std::stringstream ss;
ss << "inverse_impl: triangular inversion failed with error code " << info;

7
python/tests/test_linalg.py
View File

@ -175,6 +175,13 @@ class TestLinalg(mlx_tests.MLXTestCase):
mx.allclose(M @ M_inv, mx.eye(M.shape[0]), rtol=0, atol=1e-5)
)
# Ensure that tri_inv will 0-out the supposedly 0 triangle
x = mx.random.normal((2, 8, 8))
y1 = mx.linalg.tri_inv(x, upper=True, stream=mx.cpu)
y2 = mx.linalg.tri_inv(x, upper=False, stream=mx.cpu)
self.assertTrue(mx.all(y1 == mx.triu(y1)))
self.assertTrue(mx.all(y2 == mx.tril(y2)))
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