[CUDA] Implement SegmentedMM (#3238)

python/tests/cuda_skip.py

@@ -6,8 +6,6 @@ cuda_skip = {
     "TestBlas.test_gather_matmul",
     "TestBlas.test_gather_matmul_grad",
     "TestBlas.test_gather_mm_sorted_vjp",
-    # Segmented matmul NYI
-    "TestBlas.test_segmented_mm",
     # FFTs NYI
     "TestFFT.test_fft",
     "TestFFT.test_fft_big_powers_of_two",

python/tests/test_conv.py

@@ -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))