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