From 1da09eb82dfd5e2e7d1b579ea05a51ebb6cb6f06 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 24 Jun 2026 14:07:01 -0700 Subject: [PATCH] //4 --- tinygrad/codegen/__init__.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index c823ef4777..6af30fb721 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -59,7 +59,7 @@ pm_no_weakints = PatternMatcher([ pm_fix_image_shrink = PatternMatcher([ (UPat(Ops.SHRINK, src=(UPat(Ops.RESHAPE, src=(UPat(Ops.PARAM, name="img"), UPat())), UPat(name="idx"), UPat()), name="out"), - lambda img,out,idx: img.reshape(-1, 4).index(idx.src[0]) if len(img.shape) == 3 else None), + lambda img,out,idx: img.reshape(-1, 4).index(idx.src[0]//4) if len(img.shape) == 3 and out.shape == (4,) else None), ]) def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: @@ -133,7 +133,8 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: sink = memory_coalesing(sink, ren) # image fixup - sink = graph_rewrite(sink, pm_mops+pm_fix_image_shrink, name="fix image shrink") + if IMAGE and ren.target.device in {"QCOM", "CL", "PYTHON", "NULL"}: + sink = graph_rewrite(sink, pm_mops+pm_fix_image_shrink, name="fix image shrink") # instruction selection decompositions pm_decomp = pm_decomp+\