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

fix conv2d bug + faster conv 1d (#2195)

Browse Source 
* fix conv2d bug + faster conv 1d

* revert sort + flaky test
This commit is contained in:
Awni Hannun
2025-05-18 06:05:11 -07:00
committed by GitHub
parent 0654543dcc
commit 8576e6fe36
6 changed files with 170 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
python/tests/test_conv.py
View File

@@ -1152,6 +1152,27 @@ class TestConv(mlx_tests.MLXTestCase):
)
self.assertEqual(grads.shape, k_shape)
def test_1d_conv_with_2d(self):
x = mx.random.uniform(shape=(2, 10, 16))
y = mx.random.normal(shape=(16, 3, 16))
out = mx.conv1d(x, y, padding=1)
out_2d = mx.conv2d(
mx.expand_dims(x, axis=2), mx.expand_dims(y, axis=2), padding=(1, 0)
)
self.assertTrue(mx.allclose(out, out_2d.squeeze(2)))
x = mx.random.uniform(shape=(2, 10, 4))
y = mx.random.normal(shape=(4, 3, 4))
out = mx.conv1d(x, y, padding=1)
out_2d = mx.conv2d(
mx.expand_dims(x, axis=2), mx.expand_dims(y, axis=2), padding=(1, 0)
)
self.assertTrue(mx.allclose(out, out_2d.squeeze(2)))
if __name__ == "__main__":
unittest.main()

1
python/tests/test_vmap.py
View File

@@ -634,6 +634,7 @@ class TestVmap(mlx_tests.MLXTestCase):
self.assertEqual(fy.shape, (4, 5, 6, 7))
def test_leaks(self):
mx.synchronize()
if mx.metal.is_available():
mem_pre = mx.get_active_memory()
else: