Merge branch 'main' into nccl_backend

python/tests/cuda_skip.py

@@ -15,15 +15,6 @@ cuda_skip = {
     "TestOps.test_hadamard_grad_vmap",
     # Convolutions NYI
     "TestConv.test_1d_conv_with_2d",
-    "TestConv.test_conv_1d_groups_flipped",
-    "TestConv.test_conv_general_flip_grad",
-    "TestConv.test_conv_groups_grad",
-    "TestConv.test_torch_conv_2D",
-    "TestConv.test_torch_conv_depthwise",
-    "TestConv.test_torch_conv_general",
-    "TestConvTranspose.test_torch_conv_transpose_1D_grad",
-    "TestConvTranspose.test_torch_conv_transpose_2D_grad",
-    "TestConvTranspose.test_torch_conv_transpose_3D_grad",
     # FFTs NYI
     "TestFFT.test_fft",
     "TestFFT.test_fft_big_powers_of_two",

python/tests/test_array.py

@@ -1365,6 +1365,9 @@ class TestArray(mlx_tests.MLXTestCase):
 
     def test_array_at(self):
         a = mx.array(1)
+        with self.assertRaises(ValueError):
+            a.at.add(1)
+
         a = a.at[None].add(1)
         self.assertEqual(a.item(), 2)
 

python/tests/test_fast.py

@@ -639,12 +639,12 @@ class TestFast(mlx_tests.MLXTestCase):
             ],
             grid=(6, 1, 1),
             threadgroup=(2, 1, 1),
-            output_shapes=[(2, 2), (3, 2)],
+            output_shapes=[(3, 2), (3, 2)],
             output_dtypes=[mx.float32, mx.int32],
             stream=mx.gpu,
         )
 
-        self.assertTrue(mx.allclose(out[0], mx.full((2, 2), 14.0484)))
+        self.assertTrue(mx.allclose(out[0], mx.full((3, 2), 14.0484)))
         self.assertTrue(mx.allclose(out[1], mx.full((3, 2), -2, dtype=mx.int32)))
 
     @unittest.skipIf(not mx.metal.is_available(), "Metal is not available")