From bef5f717bc498d3374965de86b0f4070a19ee71e Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Sat, 16 May 2026 18:09:19 +0300 Subject: [PATCH] fix nolocals and beam (#16232) --- tinygrad/engine/realize.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 4c99b973c6..335617056b 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -243,9 +243,9 @@ if getenv("HCQ2"): from extra.hcq2.hcq2 import pm_hcq_exec pm_exec = pm_hcq_exec + pm_exec -def compile_linear(linear:UOp, beam=0, validate=False) -> UOp: +def compile_linear(linear:UOp, beam:int|None=None, validate=False) -> UOp: if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True) - if (beam_val:=(beam or BEAM.value)) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True) + if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True) linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True) return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True) @@ -260,4 +260,5 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None else: from tinygrad.tensor import Tensor with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False) + call = compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0).src[0] return cast(float, pm_exec.rewrite(call, ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)))