From 5bd99dd5ec43d42c6e4bef48d16d891c567aad8a Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Sun, 25 Jan 2026 16:40:57 -0800 Subject: [PATCH] Fix flaky macOS test (#3063) --- python/tests/test_blas.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/python/tests/test_blas.py b/python/tests/test_blas.py index 469d2407..dedfa5d4 100644 --- a/python/tests/test_blas.py +++ b/python/tests/test_blas.py @@ -293,6 +293,15 @@ class TestBlas(mlx_tests.MLXTestCase): out_npy = vec_npy @ mat_npy out_mlx = vec_mlx @ mat_mlx + # Due to some bug, numpy sometimes has NaNs on macOS + # See https://github.com/ml-explore/mlx/pull/3063 + nans = np.isnan(out_npy) + if np.any(nans): + nan_ids = np.where(nans) + mlx_nan_ids = tuple(mx.array(n) for n in nan_ids) + out_npy[nan_ids] = 0.0 + out_mlx[mlx_nan_ids] = 0.0 + self.assertListEqual(list(out_npy.shape), list(out_mlx.shape)) self.assertTrue(np.allclose(out_mlx, out_npy, atol=1e-5))