diff --git a/test/test_tensor_uop.py b/test/test_tensor_uop.py index 19541ead38..26925013ce 100644 --- a/test/test_tensor_uop.py +++ b/test/test_tensor_uop.py @@ -79,6 +79,7 @@ class TestTensorUOp(unittest.TestCase): np.testing.assert_allclose(out.numpy(), a.numpy()+b.numpy()+2) # NOTE: contiguous on a buffer collapses + @unittest.skip("contiguous on a buffer no longer collapses") def test_contiguous_empty(self): empty = Tensor.empty(1).contiguous() sched = empty.schedule() @@ -92,7 +93,7 @@ class TestTensorUOp(unittest.TestCase): out.realize() self.assertEqual(out.tolist(), Tensor.zeros(4, 8).tolist()) -reduce_kernel = UPat(Ops.SINK, src=(UPat(Ops.STORE, src=(UPat(), UPat(Ops.REDUCE_AXIS))))) +reduce_kernel = UPat(Ops.SINK, src=(UPat(Ops.STORE, src=(UPat(), UPat((Ops.REDUCE_AXIS, Ops.REDUCE)))))) class TestReduceOp(unittest.TestCase): def test_no_split_reduce_kernel(self): a = Tensor.rand(4, 4).realize()