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

QR factorization (#310)

Browse Source 
* add qr factorization

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
taher
2024-01-26 09:27:31 -08:00
committed by GitHub
parent 2463496471
commit 077c1ee64a
20 changed files with 322 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
python/tests/test_linalg.py
View File

@@ -89,6 +89,37 @@ class TestLinalg(mlx_tests.MLXTestCase):
out_mx = mx.linalg.norm(x_mx, ord="fro")
self.assertTrue(np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6))
def test_qr_factorization(self):
with self.assertRaises(ValueError):
mx.linalg.qr(mx.array(0.0))
with self.assertRaises(ValueError):
mx.linalg.qr(mx.array([0.0, 1.0]))
with self.assertRaises(ValueError):
mx.linalg.qr(mx.array([[0, 1], [1, 0]]))
A = mx.array([[2.0, 3.0], [1.0, 2.0]])
Q, R = mx.linalg.qr(A, stream=mx.cpu)
out = Q @ R
self.assertTrue(mx.allclose(out, A))
out = Q @ Q
self.assertTrue(mx.allclose(out, mx.eye(2), rtol=1e-5, atol=1e-7))
self.assertTrue(mx.allclose(mx.tril(R, -1), mx.zeros_like(R)))
self.assertEqual(Q.dtype, mx.float32)
self.assertEqual(R.dtype, mx.float32)
# Multiple matrices
B = mx.array([[-1.0, 2.0], [-4.0, 1.0]])
A = mx.stack([A, B])
Q, R = mx.linalg.qr(A, stream=mx.cpu)
for a, q, r in zip(A, Q, R):
out = q @ r
self.assertTrue(mx.allclose(out, a))
out = q @ q
self.assertTrue(mx.allclose(out, mx.eye(2), rtol=1e-5, atol=1e-7))
self.assertTrue(mx.allclose(mx.tril(r, -1), mx.zeros_like(r)))
if __name__ == "__main__":
unittest.main()