diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index f5e7a6af3d..1f39740eb8 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -326,7 +326,10 @@ def apply_opts(ctx:Renderer, ast:UOp): # NOTE: hand_coded_optimizations doesn't support multiblock opts yet if all(len(u.src) == 1 for u in ast.parents if u.op is Ops.LOAD): # TODO: why is the returned k from hand_coded_optimizations different than this? - for opt in hand_coded_optimizations(k).applied_opts: k.apply_opt(opt) + hc_k = hand_coded_optimizations(k) + for opt in hc_k.applied_opts: k.apply_opt(opt) + assert hc_k.applied_opts == k.applied_opts + assert hc_k.get_optimized_ast(name_override="k") == k.get_optimized_ast(name_override="k") return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None) pm_postrange_opt = PatternMatcher([