From 50b7b283dc0c83eba282bb7595fbf77fff353cd4 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 3 Feb 2026 17:41:54 +0800 Subject: [PATCH] allow after on param --- tinygrad/uop/spec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index a09a59fb0d..56e6458149 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -130,7 +130,7 @@ _tensor_spec = PatternMatcher([ (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])), # AFTER if things were kernelized - (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.AFTER)),), allow_any_len=True), lambda: True), + (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.AFTER, Ops.PARAM)),), allow_any_len=True), lambda: True), # Tensor range bind / store (UPat(Ops.BIND, (dtypes.int,dtypes.index,), (UPat(Ops.DEFINE_VAR), UPat(Ops.RANGE)), arg=None), lambda: True),