From b18aa00bbad1f3deb8fd2f3881de0001daba0b24 Mon Sep 17 00:00:00 2001
From: Roelof van Dijk <3604013+roelofvandijk@users.noreply.github.com>
Date: Fri, 12 Jul 2024 16:36:57 +0200
Subject: [PATCH] refactor: consolidate replace [run_process_replay] (#5403)

---
 tinygrad/codegen/uopgraph.py | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/tinygrad/codegen/uopgraph.py b/tinygrad/codegen/uopgraph.py
index 4991c4b7c1..0d789241aa 100644
--- a/tinygrad/codegen/uopgraph.py
+++ b/tinygrad/codegen/uopgraph.py
@@ -143,13 +143,8 @@ def expand_nodes(parents:Set[UOp], expands:List[UOp], base:UOp) -> List[UOp]:
         replace[cc] = super_replace
     else:
       for replace in replaces:
-        if cc in replace:
-          # NOTE: handle expands that are already replaced
-          tcc = replace[cc]
-          replace[cc] = UOp(tcc.op, tcc.dtype, tuple(replace.get(x, x) for x in tcc.src), tcc.arg)
-        else:
-          replace[cc] = UOp(cc.op, cc.dtype, tuple(replace.get(x, x) for x in cc.src), cc.arg)
-
+        tcc = replace.get(cc, cc)  # NOTE: handle expands that are already replaced
+        replace[cc] = UOp(tcc.op, tcc.dtype, tuple(replace.get(x, x) for x in tcc.src), tcc.arg)
   return [x.get(base, base) for x in replaces]
 
 # ***** reduce+image+contract handling *****