L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests 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 b74b2902016204117040949231887f0622bc2c39.

* Make sure that tri_inv returns a triangular matrix

---------

Co-authored-by: Angelos Katharopoulos <[email protected]>
This commit is contained in:
Abe Leininger
2025-02-19 12:42:33 -08:00
committed by GitHub
co-authored by Angelos Katharopoulos
parent 71de73a668
commit 344a29506e
2 changed files with 23 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/test_linalg.py
+7
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