zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-06 12:09:43 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Fix flaky addmm tests (#2581)

Browse Source
This commit is contained in:
Cheng
2025-09-10 14:22:22 +09:00
committed by GitHub
parent 44cc5da4bc
commit 52b8384d10
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/tests/test_blas.py
View File

@@ -702,7 +702,7 @@ class TestBlas(mlx_tests.MLXTestCase):
b = mx.ones((5, 5)) b = mx.ones((5, 5))
out = mx.addmm(a, b, a, beta=beta, alpha=alpha) out = mx.addmm(a, b, a, beta=beta, alpha=alpha)
expected = beta * a + alpha * (b @ a) expected = beta * a + alpha * (b @ a)
self.assertTrue(mx.allclose(expected, out)) self.assertTrue(mx.allclose(expected, out, atol=1e-5))
# Broadcast c # Broadcast c
a = mx.ones((5, 5)) a = mx.ones((5, 5))
@@ -710,7 +710,7 @@ class TestBlas(mlx_tests.MLXTestCase):
c = mx.ones((1, 5)) c = mx.ones((1, 5))
out = mx.addmm(c, a, b, beta=beta, alpha=alpha) out = mx.addmm(c, a, b, beta=beta, alpha=alpha)
expected = beta * c + alpha * (a @ b) expected = beta * c + alpha * (a @ b)
self.assertTrue(mx.allclose(expected, out)) self.assertTrue(mx.allclose(expected, out, atol=1e-5))
def test_addmm_grad(self): def test_addmm_grad(self):
def make_ref_addmm(alpha, beta): def make_ref_addmm(alpha, beta):