zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 00:04:41 +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
2 changed files with 23 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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