From 32df46cd7337c02e8428126c896bd686ddbeb673 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Tue, 17 Dec 2024 23:34:17 -0800
Subject: [PATCH] test const pattern [pr] (#8304)

* test const pattern [pr]

* add model to test_tiny
---
 test/test_tiny.py                           | 21 ++++++++++++++++++++-
 test/unit/test_tensor_uop_representation.py |  6 ++++++
 tinygrad/engine/schedule.py                 |  1 +
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/test/test_tiny.py b/test/test_tiny.py
index 3bd2fda82a..9001e0078c 100644
--- a/test/test_tiny.py
+++ b/test/test_tiny.py
@@ -1,6 +1,6 @@
 # basic self-contained tests of the external functionality of tinygrad
 import unittest, random
-from tinygrad import Tensor, Context, Variable, TinyJit, dtypes, Device
+from tinygrad import Tensor, Context, Variable, TinyJit, dtypes, Device, nn
 from tinygrad.helpers import IMAGE
 
 class TestTiny(unittest.TestCase):
@@ -79,6 +79,25 @@ class TestTiny(unittest.TestCase):
       ret = Tensor.ones(s).contiguous().reshape(i.bind(s)).sum()
       self.assertEqual(ret.item(), s)
 
+  # *** a model ***
+
+  def test_mnist_model(self):
+    layers = [
+      nn.Conv2d(1, 32, 5), Tensor.relu,
+      nn.Conv2d(32, 32, 5), Tensor.relu,
+      nn.BatchNorm(32), Tensor.max_pool2d,
+      nn.Conv2d(32, 64, 3), Tensor.relu,
+      nn.Conv2d(64, 64, 3), Tensor.relu,
+      nn.BatchNorm(64), Tensor.max_pool2d,
+      lambda x: x.flatten(1), nn.Linear(576, 10)]
+
+    # pre-realize random weights
+    for p in nn.state.get_parameters(layers): p.realize()
+
+    # run model inference
+    probs = Tensor.rand(1, 1, 28, 28).sequential(layers).tolist()
+    self.assertEqual(len(probs[0]), 10)
+
   # *** image ***
 
   @unittest.skipIf(Device.DEFAULT != "GPU", "image only supported on GPU")
diff --git a/test/unit/test_tensor_uop_representation.py b/test/unit/test_tensor_uop_representation.py
index 9d71eac66f..f4011fd2d2 100644
--- a/test/unit/test_tensor_uop_representation.py
+++ b/test/unit/test_tensor_uop_representation.py
@@ -3,6 +3,7 @@ from tinygrad import Tensor
 from tinygrad.ops import UPat, Ops
 
 realized_pattern = UPat(Ops.VIEW, src=(UPat(Ops.BUFFER),))
+const_pattern = UPat(Ops.VIEW, src=(UPat(Ops.BUFFER), UPat(Ops.CONST)))
 def is_pattern(ten:Tensor, pat:UPat): assert pat.match(ten.lazydata, {})
 
 class TestTensorUopRepresentation(unittest.TestCase):
@@ -18,6 +19,11 @@ class TestTensorUopRepresentation(unittest.TestCase):
     print(c.lazydata)
     is_pattern(c, UPat(Ops.ADD, src=(realized_pattern, realized_pattern)))
 
+  def test_const_pattern(self):
+    a = Tensor(1)
+    print(a.lazydata)
+    is_pattern(a, const_pattern)
+
   def test_consts_do_not_realize(self):
     a = Tensor(1)
     print(a.lazydata)
diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py
index a142a83152..e6653c4080 100644
--- a/tinygrad/engine/schedule.py
+++ b/tinygrad/engine/schedule.py
@@ -150,6 +150,7 @@ def _append_st_vars(ctx:ScheduleItemContext, x:UOp) -> Optional[UOp]:
   return st.to_uop() if st != x.st else None
 
 def _append_buf(ctx:ScheduleItemContext, x:UOp) -> UOp:
+  assert x.arg[0] != -1, "fake -1 BUFFERS should not make it here"
   ctx.bufs.append(x)
   return UOp(Ops.DEFINE_GLOBAL, x.dtype.ptr(), (), len(ctx.bufs)-1)
 append_bufs = PatternMatcher([(UPat(Ops.BUFFER, name="x"), _append_buf)])