testing backward

test/test_tiny.py

@@ -92,19 +92,23 @@ class TestTiny(unittest.TestCase):
   def test_mnist(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)]
+      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)]
 
     # replace random weights with ones
     Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])
 
     # run model inference
     probs = Tensor.empty(4, 1, 28, 28).sequential(layers).tolist()
-    self.assertEqual(len(probs[0]), 10)
+    #self.assertEqual(len(probs[0]), 10)
+
+    for x in nn.state.get_parameters(layers): x.requires_grad_()
+    Tensor.empty(4, 1, 28, 28).sequential(layers).sum().backward()
+    Tensor.realize(*[x.grad for x in nn.state.get_parameters(layers) if x.grad is not None])
 
   # *** image ***
 

tinygrad/schedule/rangeify.py

@@ -229,6 +229,7 @@ def index_child(ctx:RangeifyContext, c:UOp, x:UOp, idx:UOp):
     idx_ranges, end_ranges = ctx.seen_child[c]
     for i,nr in zip(idx_ranges, end_ranges): out_rngs[i] = nr
   if len(idx_ranges) == 0: return c.index(*out_rngs)
+  # NOTE: partial contigs can still come from here
   return c.index(*out_rngs).bufferize(*end_ranges, arg=x.device).index(*[idx.src[1+i] for i in idx_ranges])
 
 def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp):

tinygrad/tensor.py

@@ -252,7 +252,8 @@ class Tensor(MathTrait):
     # create the schedule
     schedule, var_vals = create_schedule_with_vars(sink)
     schedule = memory_planner(schedule)
-    if DEBUG >= 1 and len(schedule) >= 10: print(f"scheduled {len(schedule)} kernels in {(time.perf_counter()-st)*1000:.2f} ms")
+    if (DEBUG >= 1 and len(schedule) >= 10) or (DEBUG >= 2 and len(schedule) > 1):
+      print(f"scheduled {len(schedule)} kernels in {(time.perf_counter()-st)*1000:.2f} ms")
     return schedule, var_vals
 
   def schedule(self, *lst:Tensor) -> list[ScheduleItem]: