From bdb421f13edf01c5040c434cd8b3a1569edda9de Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Sat, 3 Jan 2026 23:09:39 -0500 Subject: [PATCH] process_replay: passthrough sink arg for Ops.PROGRAM input (#14000) --- test/external/process_replay/process_replay.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/external/process_replay/process_replay.py b/test/external/process_replay/process_replay.py index 2ebb8c17c2..15fa660ced 100755 --- a/test/external/process_replay/process_replay.py +++ b/test/external/process_replay/process_replay.py @@ -54,7 +54,7 @@ def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]: # the ast.arg is non None if we are inside of search.py sink_arg = ast.arg or KernelInfo(opts_to_apply=tuple(opts) if opts is not None else p.applied_opts if BEAM>=1 else None) - input_ast = ast.replace(arg=replace(sink_arg, name=p.name)) + input_ast = ast if ast.op is Ops.PROGRAM else ast.replace(arg=replace(sink_arg, name=p.name)) p2 = get_program(input_ast, renderer=renderer) def to_str(ret:ProgramSpec) -> str: # PYTHON renderer pickles UOps, first unpickle and decode here