Further reduction tuning (#1349)

* More reduction tuning
* Forgotten pdb
* Small column long row specialization
This commit is contained in:
Angelos Katharopoulos
2024-08-23 10:35:25 -07:00
committed by GitHub
parent da8deb2b62
commit b57a52813b
4 changed files with 77 additions and 29 deletions

View File

@@ -43,10 +43,6 @@ class TestReduce(mlx_tests.MLXTestCase):
z_npy = np.sum(y_npy, axis=a) / 1000
z_mlx = mx.sum(y_mlx, axis=a) / 1000
mx.eval(z_mlx)
if not np.allclose(z_npy, np.array(z_mlx), atol=1e-4):
import pdb
pdb.set_trace()
self.assertTrue(
np.allclose(z_npy, np.array(z_mlx), atol=1e-4)
)