diff --git a/test/test_tiny.py b/test/test_tiny.py index 4a77afdc25..1c18e537d5 100644 --- a/test/test_tiny.py +++ b/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 *** diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 7b57e67fc0..5a342ab802 100644 --- a/tinygrad/schedule/rangeify.py +++ b/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): diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index cf0c18ca68..4ce2a530f6 100644 --- a/tinygrad/tensor.py +++ b/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]: