From a8dca47fbcf47c3bef7d6a31d970d7ea650f6639 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 3 Oct 2025 10:52:18 +0800 Subject: [PATCH] fix --- tinygrad/codegen/__init__.py | 2 +- tinygrad/codegen/late/expander.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 516a128a3d..a1e6afaef8 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -19,7 +19,7 @@ from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, blo from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops from tinygrad.codegen.opt.postrange import pm_postrange_opt from tinygrad.codegen.simplify import pm_simplify_ranges, pm_reduce_simplify, pm_flatten_range, pm_split_ranges -from tinygrad.schedule.rangeify import pm_add_buffers, rangeify_codegen, pm_fix_bufferize +from tinygrad.schedule.rangeify import pm_add_buffers, rangeify_codegen @dataclass class RewriteStep: diff --git a/tinygrad/codegen/late/expander.py b/tinygrad/codegen/late/expander.py index d2d0a41162..9a42d414ce 100644 --- a/tinygrad/codegen/late/expander.py +++ b/tinygrad/codegen/late/expander.py @@ -157,6 +157,9 @@ pm_pre_expander = PatternMatcher([ # fix REDUCEs with UNROLLs (UPat(Ops.REDUCE, name="x"), fix_reduce_unroll), (UPat(Ops.STORE, name="x"), fix_store_unroll), +]) + +pm_group_for_reduce = PatternMatcher([ # fix group for reduce (UPat(Ops.REDUCE, name="x"), fix_group_for_reduce), ])