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

Check edge case handling in row reduce med kernel (#858)

Browse Source
This commit is contained in:
Jagrit Digani
2024-03-20 11:37:58 -07:00
committed by GitHub
parent cec8661113
commit b219d12a6b
2 changed files with 20 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/tests/test_ops.py
View File

@@ -590,6 +590,18 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertListEqual(list(sum_npy.shape), list(sum_mlx.shape))
self.assertTrue(np.all(sum_npy == sum_mlx))
x_npy = (
np.arange(3 * 2 * 3 * 3 * 3 * 3)
.reshape(3, 2, 3, 3, 3, 3)
.astype(np.float32)
)
x_mlx = mx.array(x_npy)
y_mlx = x_mlx.sum(axis=(0, 1, 3, 4, 5))
y_npy = x_npy.sum(axis=(0, 1, 3, 4, 5))
self.assertTrue(np.array_equal(y_mlx, y_npy))
def test_prod(self):
x = mx.array(
[