diff --git a/test/test_schedule.py b/test/test_schedule.py index b8acda4675..3eb208fbdf 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -2437,29 +2437,6 @@ class TestCopyFolding(unittest.TestCase): b.realize() self.assertListEqual(b.tolist(), [[0, 2], [1, 3]]) -class TestTensorUOpSpec(unittest.TestCase): - def test_const_must_be_unmasked(self): - a = Tensor.ones((4, 4)).pad((2, 2)) - unsafe_push_views = PatternMatcher([ - (UPat.cvar("root").view(name="view"), lambda root,view: root.replace(src=tuple(x.view(view.st) for x in root.src))), - ]) - a.uop = graph_rewrite(a.uop.sink(), merge_views+merge_views+unsafe_push_views) - with self.assertRaisesRegex(RuntimeError, "UOp verification failed"): - a.schedule() - - def test_expanded_const_ok(self): - a = Tensor.ones((4, 4)) - t = graph_rewrite(a.uop.sink(), merge_views+merge_views) - create_schedule_with_vars(t) - - # NOTE: changing symbolic CONST VIEWs is not allowed - @unittest.expectedFailure - def test_symbolic_shape_ok(self): - a = Tensor.ones(4) - vi = UOp.variable("i", 1, 10).bind(4) - a.uop = graph_rewrite(a.reshape(vi).sum().uop, merge_views+merge_views) - a.schedule() - class TestBufferUOp(unittest.TestCase): # BUFFER has a ShapeTracker of shape=(n,) and stride=(1,) def test_buffer_has_buffer(self): diff --git a/test/unit/test_verify_ast.py b/test/unit/test_uop_spec.py similarity index 93% rename from test/unit/test_verify_ast.py rename to test/unit/test_uop_spec.py index fd8a921260..dfc9bca106 100644 --- a/test/unit/test_verify_ast.py +++ b/test/unit/test_uop_spec.py @@ -5,7 +5,7 @@ from tinygrad import Tensor from tinygrad.codegen.kernel import Kernel from tinygrad.helpers import DEBUG from tinygrad.uop.ops import UOp, Ops, print_uops -from tinygrad.uop.spec import type_verify, ast_spec +from tinygrad.uop.spec import type_verify, ast_spec, tensor_uop_spec from tinygrad.shape.shapetracker import ShapeTracker from tinygrad import dtypes from tinygrad.shape.view import View @@ -23,7 +23,7 @@ def helper_test_verify_ast(*stores:UOp) -> Kernel: if DEBUG >= 4: print(k.to_program().src) return k -class TestVerifyAST(unittest.TestCase): +class TestUOpSpec(unittest.TestCase): def test_tiny_add(self): dtype = dtypes.int buf_0 = UOp(Ops.DEFINE_GLOBAL, dtype.ptr(), (), 0) @@ -94,5 +94,11 @@ class TestVerifyAST(unittest.TestCase): st = UOp.store(buf.view(ShapeTracker.from_shape(())), a.cast(dtypes.float)) helper_test_verify_ast(st) + def test_assert_masked_view_in_const(self): + t = Tensor(6).uop + a = t.replace(src=(t.src[0].replace(arg=t.st.reshape((1,)).pad(((0, 1),))),)) + with self.assertRaisesRegex(RuntimeError, "UOp verification failed"): + type_verify([a], tensor_uop_spec) + if __name__ == '__main__': unittest.main()