diff --git a/python/tests/cuda_skip.py b/python/tests/cuda_skip.py index fce92bacb..5b5ed6abb 100644 --- a/python/tests/cuda_skip.py +++ b/python/tests/cuda_skip.py @@ -3,6 +3,8 @@ cuda_skip = { "TestLayers.test_quantized_embedding", "TestOps.test_dynamic_slicing", "TestReduce.test_dtypes", + "TestReduce.test_nanpropagation", + "TestReduce.test_nanpropagation_complex64", # Block masked matmul NYI "TestBlas.test_block_masked_matmul", # Gather matmul NYI diff --git a/python/tests/test_reduce.py b/python/tests/test_reduce.py index 0342173dd..d757f1527 100644 --- a/python/tests/test_reduce.py +++ b/python/tests/test_reduce.py @@ -170,7 +170,7 @@ class TestReduce(mlx_tests.MLXTestCase): x = (mx.random.normal((4, 4)) * 10).astype(getattr(mx, dtype)) indices = mx.random.randint(0, 4, shape=(6,)).reshape(3, 2) for idx in indices: - x[*idx] = mx.nan + x[idx[0], idx[1]] = mx.nan x_np = np.array(x) for op in ["max"]: