[CUDA] Implement SegmentedMM (#3238)

2026-03-11 13:31:43 -07:00
parent 1c2d7041ab
commit a9573f92f6
7 changed files with 399 additions and 4 deletions
@@ -1151,7 +1151,7 @@ class TestConv(mlx_tests.MLXTestCase):
            )
            self.assertEqual(grads.shape, k_shape)

-    def test_1d_conv_with_2d(self):
+    def test_conv_1d_with_2d(self):
        x = mx.random.uniform(shape=(2, 10, 16))
        y = mx.random.normal(shape=(16, 3, 16))