diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 475705c263..7aaaccdce4 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -208,7 +208,7 @@ spec_program = PatternMatcher([ lambda x: False if x.dtype.count > 1 and (x.dtype.count,) != x.shape else None), # STACK/GEP in program. TODO: this should match Tensor - (UPat(Ops.STACK, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)), + (UPat(Ops.STACK, name="x"), lambda x: len(x.src)>1), (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()), # if has a