From 6122b3c98fbc1bbc840c95e095bbab9e9150d8eb Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:17:29 -0700 Subject: [PATCH] use check_schedule in tests where possible (#17400) --- test/backend/test_arange.py | 14 +++----- test/backend/test_linearizer.py | 5 ++- test/backend/test_multitensor.py | 5 ++- test/backend/test_nn.py | 10 ++---- test/backend/test_renderer_failures.py | 4 +-- test/backend/test_schedule.py | 33 +++---------------- test/backend/test_softmax_fusion.py | 4 +-- test/backend/test_transcendental.py | 2 +- test/helpers.py | 25 ++++++++++++++- test/null/test_attention.py | 5 ++- test/null/test_real_world.py | 7 ++-- test/null/test_schedule.py | 44 ++++++-------------------- test/null/test_winograd.py | 3 +- 13 files changed, 63 insertions(+), 98 deletions(-) diff --git a/test/backend/test_arange.py b/test/backend/test_arange.py index b2faa47ef9..676116f431 100644 --- a/test/backend/test_arange.py +++ b/test/backend/test_arange.py @@ -4,7 +4,7 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable from tinygrad.helpers import Context, getenv, DEV from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear from tinygrad.renderer.ptx import PTXRenderer -from test.helpers import needs_second_gpu +from test.helpers import needs_second_gpu, check_schedule class TestArange(unittest.TestCase): def _get_flops(self, tensor, desired): @@ -55,8 +55,7 @@ class TestIndexing(unittest.TestCase): with Context(NOOPT=1): GlobalCounters.reset() out = ((Tensor.arange(1,16385)-1)*needle).sum() - linear, var_vals = out.linear_with_vars() - self.assertEqual(len(linear.src), 1) + linear, var_vals = check_schedule(out, 1) run_linear(linear, var_vals) self.assertEqual(out.item(), 1337) @@ -72,8 +71,7 @@ class TestIndexing(unittest.TestCase): reshape_dataset = dataset.T.reshape(1, DDIM, DSET, 1).expand(4, DDIM, DSET, 1) full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1, buffer=False)) X = full.sum(axis=(2,3)) - linear, var_vals = X.linear_with_vars() - self.assertEqual(len(linear.src), 1) + linear, var_vals = check_schedule(X, 1) run_linear(linear, var_vals) assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}" np.testing.assert_allclose(real_index, X.numpy()) @@ -98,8 +96,7 @@ class TestIndexing(unittest.TestCase): GlobalCounters.reset() X = dataset[idxs] assert X.shape == (4,DDIM) - linear, var_vals = X.linear_with_vars() - self.assertEqual(len(linear.src), 1) + linear, var_vals = check_schedule(X, 1) run_linear(linear, var_vals) assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}" np.testing.assert_allclose(real_index, X.numpy()) @@ -113,8 +110,7 @@ class TestIndexing(unittest.TestCase): GlobalCounters.reset() X = dataset[idxs] assert X.shape == (4,DDIM) - linear, var_vals = X.linear_with_vars() - self.assertEqual(len(linear.src), 1) + linear, var_vals = check_schedule(X, 1) run_linear(linear, var_vals) assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}" np.testing.assert_allclose(real_index, X.numpy()) diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py index 6d833cbeba..958a9be840 100644 --- a/test/backend/test_linearizer.py +++ b/test/backend/test_linearizer.py @@ -12,7 +12,7 @@ from tinygrad.dtype import DType, dtypes, AddrSpace from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.cstyle import CUDARenderer from tinygrad.renderer.isa import ISARenderer -from test.helpers import replace_opts +from test.helpers import replace_opts, check_schedule from test.backend.test_softmax_fusion import single_kernel_softmax MOCKGPU = DEV.interface.startswith("MOCK") @@ -293,8 +293,7 @@ class TestLinearizer(unittest.TestCase): a = Tensor.ones(4, 4).contiguous().realize() b = a.shrink(((1, 2), None)).pad(((1, 2), None)).bool() a.assign(b.where(2, a)) - linear, var_vals = a.linear_with_vars() - assert len(linear.src) == 1 + linear, var_vals = check_schedule(a, 1) run_linear(linear, var_vals) np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.]) program = to_program(replace_opts(linear.src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer) diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index 5e461236c1..7ef0a6888b 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -6,7 +6,7 @@ from tinygrad.nn.state import get_parameters from tinygrad.engine.realize import run_linear, compile_linear import numpy as np from hypothesis import given, strategies as strat, settings -from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph +from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False)) settings.load_profile("my_profile") @@ -355,8 +355,7 @@ class TestMultiTensor(unittest.TestCase): def test_const_like_shrink_on_shard_axis(self): t = Tensor.ones(16, 16, dtype=dtypes.int).shard(devices_2, axis=0) out = t.const_like(2)[:, :8] - linear, var_vals = out.linear_with_vars() - self.assertEqual(len(linear.src), 0) + linear, var_vals = check_schedule(out, 0) run_linear(linear, var_vals) self.assertEqual(out.tolist(), [[2]*8]*16) diff --git a/test/backend/test_nn.py b/test/backend/test_nn.py index 88c4174c4f..75195d88e0 100644 --- a/test/backend/test_nn.py +++ b/test/backend/test_nn.py @@ -3,11 +3,11 @@ import unittest import numpy as np import torch from tinygrad import Tensor, Device, TinyJit, dtypes -from tinygrad.uop.ops import Ops from tinygrad.helpers import GlobalCounters, Context from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell from tinygrad.nn.state import load_state_dict +from test.helpers import check_schedule from tinygrad.engine.realize import run_linear from test.helpers import not_support_multi_device, needs_second_gpu, slow @@ -428,18 +428,14 @@ class TestNN(unittest.TestCase): a = Tensor([[1, 5, 9, 11], [12, 19, 8, 1]]) result = layer(a) - linear, var_vals = result.linear_with_vars() - self.assertEqual(len([call for call in linear.src if call.src[0].op is Ops.SINK]), kcount, - "first run realizes weight and embedding") + linear, var_vals = check_schedule(result, kcount) run_linear(linear, var_vals) b = Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) result = layer(b) - linear, var_vals = result.linear_with_vars() - self.assertEqual(1, len([call for call in linear.src if call.src[0].op is Ops.SINK]), - "second run realizes embedding only") + linear, var_vals = check_schedule(result, 1) run_linear(linear, var_vals) print(f"Embedding used {GlobalCounters.global_ops} ops") self.assertLessEqual(GlobalCounters.global_ops, ops) diff --git a/test/backend/test_renderer_failures.py b/test/backend/test_renderer_failures.py index 826928bece..d1f7a145fa 100644 --- a/test/backend/test_renderer_failures.py +++ b/test/backend/test_renderer_failures.py @@ -8,6 +8,7 @@ from tinygrad.helpers import prod from tinygrad.renderer.cstyle import CStyleLanguage from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.wgsl import WGSLRenderer +from test.helpers import check_schedule from tinygrad.runtime.ops_python import PythonRenderer from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu from tinygrad.tensor import Tensor @@ -61,8 +62,7 @@ class TestCStyleFailures(unittest.TestCase): dtype = "bool" if op in (Ops.OR, Ops.XOR, Ops.AND) else None ret = Tensor.empty(1, dtype=dtype) for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype)) - linear = ret.schedule_linear() - assert len(linear.src) == 1 + linear, _ = check_schedule(ret, 1) src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[2].arg self.assertEqual("("*5 not in src, should_strip_paren) diff --git a/test/backend/test_schedule.py b/test/backend/test_schedule.py index 838b8cbcdd..2d86674633 100644 --- a/test/backend/test_schedule.py +++ b/test/backend/test_schedule.py @@ -6,34 +6,13 @@ import unittest, time import numpy as np from tinygrad import nn, dtypes, Device, Tensor, Variable -from tinygrad.uop.ops import UOp, Ops, UPat -from tinygrad.helpers import DEBUG, DEV, GlobalCounters, Context, all_same, temp -from tinygrad.engine.realize import compile_linear, run_linear +from tinygrad.uop.ops import Ops, UPat +from tinygrad.helpers import DEV, GlobalCounters, Context, all_same, temp +from tinygrad.engine.realize import run_linear +from test.helpers import check_schedule supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes() -class KernelCountException(Exception): pass -def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True): - if to_prerealize: - with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize) - if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars() - elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t) - else: - assert isinstance(t, UOp), f"can't schedule {t}" - linear, var_vals = Tensor(t).linear_with_vars() - kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1) - for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink) - if kernel_cnt != allowed: - print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}") - if DEBUG >= 3: - for i,call in enumerate(linear.src): - print("kernel", i+1) - print(call.src[0]) - raise KernelCountException(f"{kernel_cnt} != {allowed}") - # test compiling the linear - compile_linear(linear) - return linear, var_vals - def _realize_weights(m): for p in nn.state.get_parameters(m): p.realize() @@ -113,11 +92,9 @@ class TestSchedule(unittest.TestCase): a2 = mop(a) expected = (a+a2).tolist() a.assign(a+a2) - linear, var_vals = a.linear_with_vars() - kcount = len(linear.src) + linear, var_vals = check_schedule(a, expected_kcount) run_linear(linear, var_vals) self.assertListEqual(a.tolist(), expected) - self.assertEqual(kcount, expected_kcount) def test_setitem_permuted_sched(self): self.test_setitem_sched(lambda x: x.T, 2) def test_setitem_paddded_sched(self): self.test_setitem_sched(lambda x: x.shrink_to(4, 1).pad_to(4, 4), 1) diff --git a/test/backend/test_softmax_fusion.py b/test/backend/test_softmax_fusion.py index eb0dda3ade..97f76139ee 100644 --- a/test/backend/test_softmax_fusion.py +++ b/test/backend/test_softmax_fusion.py @@ -4,6 +4,7 @@ from tinygrad import Tensor, GlobalCounters, Context, Device from tinygrad.dtype import DTypeLike, dtypes from tinygrad.engine.realize import run_linear from tinygrad.helpers import DEBUG, get_single_element +from test.helpers import check_schedule def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Tensor: # only support axis =-1 @@ -103,8 +104,7 @@ class TestFuse(unittest.TestCase): k = (x @ wk).contiguous() v = (x @ wv).contiguous() attn = q.scaled_dot_product_attention(k, v) - s = attn.schedule_linear() - self.assertEqual(len(s.src), 4) # 3 matmul and 1 attention + check_schedule(attn, 4) # 3 matmul and 1 attention @unittest.skip("needs RANGEIFY>1") def test_flash_attention(self): diff --git a/test/backend/test_transcendental.py b/test/backend/test_transcendental.py index fb96ad5147..20f008154d 100644 --- a/test/backend/test_transcendental.py +++ b/test/backend/test_transcendental.py @@ -2,7 +2,7 @@ import unittest from tinygrad import Tensor, Device, dtypes from tinygrad.tensor import _to_np_dtype from tinygrad.helpers import Context, getenv, DEV, OSX -from test.backend.test_schedule import check_schedule +from test.helpers import check_schedule from test.backend.test_dtype_alu import ht, dtypes_float import numpy as np import math diff --git a/test/helpers.py b/test/helpers.py index 1e2b893d16..c613d78f1f 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -8,10 +8,11 @@ from tinygrad.tensor import _to_np_dtype from tinygrad.codegen import to_program from tinygrad.dtype import DType, truncate from tinygrad.nn.state import get_parameters -from tinygrad.helpers import T, Target, DEV +from tinygrad.helpers import T, Target, DEV, DEBUG, Context from tinygrad.renderer import Renderer from tinygrad.codegen import full_rewrite_to_sink, line_rewrite, pm_linearize_cleanups from tinygrad.codegen.late.linearizer import linearize +from tinygrad.engine.realize import compile_linear # decorator to skip slow tests by default, run with RUN_SLOW=1 to include them slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run") @@ -34,6 +35,28 @@ def derandomize_model(model): p.replace(Tensor.empty(p.shape, device=p.device, dtype=p.dtype)) p.realize() +class KernelCountException(Exception): pass +def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True): + if to_prerealize: + with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize) + if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars() + elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t) + else: + assert isinstance(t, UOp), f"can't schedule {t}" + linear, var_vals = Tensor(t).linear_with_vars() + kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1) + for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink) + if kernel_cnt != allowed: + print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}") + if DEBUG >= 3: + for i,call in enumerate(linear.src): + print("kernel", i+1) + print(call.src[0]) + raise KernelCountException(f"{kernel_cnt} != {allowed}") + # test compiling the linear + compile_linear(linear) + return linear, var_vals + def call_is_graph(call:UOp) -> bool: ast = call.src[0] return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph" diff --git a/test/null/test_attention.py b/test/null/test_attention.py index aa23608de9..2c03e58450 100644 --- a/test/null/test_attention.py +++ b/test/null/test_attention.py @@ -1,7 +1,7 @@ import unittest from tinygrad import Tensor, dtypes, TinyJit, UOp from tinygrad.llm.model import apply_rope as apply_rope_new, precompute_freqs_cis -from test.helpers import assert_jit_cache_len +from test.helpers import assert_jit_cache_len, check_schedule def apply_rope(x:Tensor, start_pos:int): B, H, T, Hd = x.shape @@ -16,9 +16,8 @@ class TestAttention(unittest.TestCase): k = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize() v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize() attn = q.scaled_dot_product_attention(k, v) - sched = attn.schedule_linear() # attention has 4 kernels now - self.assertEqual(len(sched.src), 4) + check_schedule(attn, 4) def test_apply_rope_jit_prune(self): def rope_fn(x_in, pos): return apply_rope(x_in, pos) diff --git a/test/null/test_real_world.py b/test/null/test_real_world.py index 378e8b924b..23dcab609f 100644 --- a/test/null/test_real_world.py +++ b/test/null/test_real_world.py @@ -5,7 +5,7 @@ from tinygrad.nn.state import get_parameters from tinygrad.engine.jit import TinyJit from tinygrad import Tensor, Device, GlobalCounters, dtypes, Variable from tinygrad.helpers import Context -from test.helpers import slow, jit_cache_count +from test.helpers import slow, jit_cache_count, KernelCountException from extra.lr_scheduler import OneCycleLR from test.helpers import derandomize_model @@ -35,8 +35,9 @@ def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jit assert mem_used < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.3f} GB - {mem_used:.3} GB used" assert (max_memory_allowed - mem_used) / max_memory_allowed < 0.2, f"{max_memory_allowed:.3f} GB is too far from {mem_used:.3} GB used" if kernels_used: - assert kernels_used <= max_kernels_allowed, f"{nm} used more than {max_kernels_allowed} kernels, it used {kernels_used}" - assert (max_kernels_allowed - kernels_used) / max_kernels_allowed < 0.2, f"{max_kernels_allowed=} is too far from {kernels_used=} used" + if kernels_used > max_kernels_allowed: raise KernelCountException(f"{nm} used more than {max_kernels_allowed} kernels, it used {kernels_used}") + if (max_kernels_allowed - kernels_used) / max_kernels_allowed >= 0.2: + raise KernelCountException(f"{max_kernels_allowed=} is too far from {kernels_used=} used") if all_jitted: assert kernels_used > 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted" # noqa: E501 diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index 2f4cbae50a..29dd939281 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -3,31 +3,10 @@ import gc, unittest, time from typing import cast from tinygrad import nn, dtypes, Device, Tensor, getenv from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo -from tinygrad.helpers import DEBUG, GlobalCounters, Context -from tinygrad.engine.realize import compile_linear, run_linear +from tinygrad.helpers import GlobalCounters, Context +from tinygrad.engine.realize import run_linear, compile_linear from tinygrad.codegen import to_program - -class KernelCountException(Exception): pass -def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True): - if to_prerealize: - with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize) - if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars() - elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t) - else: - assert isinstance(t, UOp), f"can't schedule {t}" - linear, var_vals = Tensor(t).linear_with_vars() - kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1) - for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink) - if kernel_cnt != allowed: - print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}") - if DEBUG >= 3: - for i,call in enumerate(linear.src): - print("kernel", i+1) - print(call.src[0]) - raise KernelCountException(f"{kernel_cnt} != {allowed}") - # test compiling the linear - compile_linear(linear) - return linear, var_vals +from test.helpers import check_schedule def _realize_weights(m): for p in nn.state.get_parameters(m): p.realize() @@ -143,7 +122,7 @@ class TestSimpleSchedule(unittest.TestCase): a = Tensor.empty(16,16).sum(axis=1) a1 = a.reshape(4,4) a2 = a.reshape(16,1,1) - self.assertEqual(len(Tensor.schedule_linear(a1, a2).src), 1) + check_schedule([a1, a2], 1) class TestSchedule(unittest.TestCase): def setUp(self): @@ -155,8 +134,7 @@ class TestSchedule(unittest.TestCase): def test_arange_avgpool2d(self, kcount=1): x = Tensor.arange(25).reshape(1,1,5,5).cast(dtypes.float32) t = x.avg_pool2d(padding=1).clone() - linear, var_vals = t.linear_with_vars() - self.assertEqual(len(linear.src), kcount) + check_schedule(t, kcount) def test_arange_avgpool2d_fused_noopt(self): with Context(NOOPT=1): self.test_arange_avgpool2d(kcount=1) @@ -874,8 +852,7 @@ class TestSchedule(unittest.TestCase): t = Tensor.zeros((3, 3)).contiguous().realize() v = t[1] # view - is_realized but not has_buffer_identity assert v.uop.is_realized - linear, _ = Tensor.linear_with_vars(v) - self.assertEqual(len(linear.src), 0) + check_schedule(v, 0) # NOTE: because empty does not have a lowered kernel if realize is called on a childless empty, it never gets allocated. def test_childless_empty_never_allocates(self): @@ -1457,8 +1434,7 @@ class TestSchedule(unittest.TestCase): Tensor.manual_seed(0) x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.half).realize() out = x.softmax(dtype=dtypes.float) - linear = out.schedule_linear() - self.assertEqual(len(linear.src), 3) + linear, _ = check_schedule(out, 3) # max reduction stays in input dtype (no numerical loss), upcast happens after subtracting max self.assertEqual(linear.src[0].src[1].dtype, dtypes.half) self.assertEqual(linear.src[1].src[1].dtype, dtypes.float) @@ -1873,8 +1849,7 @@ class TestFusionOp(unittest.TestCase): val = 1.0 a = Tensor(val) for _ in range(24): a = Tensor.stack(a, a)[0] - linear = a.schedule_linear() - self.assertLessEqual(len(linear.src), 1) + check_schedule(a, 0) self.assertLess(time.perf_counter()-st, 2.0) def test_recursive_reshape(self): @@ -1883,8 +1858,7 @@ class TestFusionOp(unittest.TestCase): b = Tensor.empty(16, 2).realize() r = a.sum(1) for _ in range(24): r = r.reshape(16, 2) + b - linear = r.schedule_linear() - self.assertEqual(len(linear.src), 1) + check_schedule(r, 1) self.assertLess(time.perf_counter()-st, 2.0) # NOTE: the NULL backend supports SLICE diff --git a/test/null/test_winograd.py b/test/null/test_winograd.py index 1530675828..0fa0e5039e 100644 --- a/test/null/test_winograd.py +++ b/test/null/test_winograd.py @@ -1,6 +1,7 @@ import unittest, sys from tinygrad import Tensor, GlobalCounters, dtypes, Context from tinygrad.helpers import WINO +from test.helpers import check_schedule @unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows") class TestWinograd(unittest.TestCase): @@ -13,7 +14,7 @@ class TestWinograd(unittest.TestCase): def test_forward_kernels(self): x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize() out = Tensor.conv2d(x,w) - self.assertEqual(len(out.schedule_linear().src), 4) + check_schedule(out, 4) def test_backward_counters(self): # contiguous_backward on the pooled input keeps the input-transform adjoint out of the overlap accumulation, so