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

Allow non-square lu (#1889)

Browse Source
This commit is contained in:
Awni Hannun
2025-02-20 08:13:23 -08:00
committed by GitHub
parent c86422bdd4
commit bbda0fdbdb
3 changed files with 34 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/tests/test_linalg.py
View File

@@ -358,6 +358,15 @@ class TestLinalg(mlx_tests.MLXTestCase):
L = mx.take_along_axis(L, P[..., None], axis=-2)
self.assertTrue(mx.allclose(L @ U, a))
# Test non-square matrix
a = mx.array([[3.0, 1.0, 2.0], [1.0, 8.0, 6.0]])
P, L, U = mx.linalg.lu(a, stream=mx.cpu)
self.assertTrue(mx.allclose(L[P, :] @ U, a))
a = mx.array([[3.0, 1.0], [1.0, 8.0], [9.0, 2.0]])
P, L, U = mx.linalg.lu(a, stream=mx.cpu)
self.assertTrue(mx.allclose(L[P, :] @ U, a))
def test_lu_factor(self):
mx.random.seed(7)