map is only the map

test/unit/test_rewrite_map.py

@@ -11,12 +11,9 @@ class TestRewriteMap(unittest.TestCase):
     root_add = zero_node + inner_add
 
     # Perform top-down rewrite
-    rewritten_sink, node_map = graph_rewrite_map(root_add, symbolic)
+    node_map = graph_rewrite_map(root_add, symbolic)
 
     # We expect that add(0, add(0, 5)) -> add(0, 5) -> 5
-    # So the final node should be 'five_node'
-    assert rewritten_sink == five_node
-
     # Check the mapping
     assert node_map[root_add] == five_node
     assert node_map[inner_add] == five_node
@@ -36,10 +33,8 @@ class TestRewriteMap(unittest.TestCase):
     neg_five = -five_node
     double_neg_five = -neg_five
 
-    rewritten_sink, node_map = graph_rewrite_map(double_neg_five, symbolic)
+    node_map = graph_rewrite_map(double_neg_five, symbolic)
 
-    # Expect neg(neg(5)) -> 5
-    self.assertEqual(rewritten_sink, five_node)
     # node_map should map double_neg_five -> five_node
     self.assertEqual(node_map[double_neg_five], five_node)
     # five_node maps to itself
@@ -55,10 +50,8 @@ class TestRewriteMap(unittest.TestCase):
     double_neg_five = -neg_five
     root_add = zero_node + double_neg_five
 
-    rewritten_sink, node_map = graph_rewrite_map(root_add, symbolic)
+    node_map = graph_rewrite_map(root_add, symbolic)
 
-    # Expect final node is 'five_node'
-    self.assertEqual(rewritten_sink, five_node)
     # node_map: root_add -> five_node, double_neg_five -> five_node
     self.assertEqual(node_map[root_add], five_node)
     self.assertEqual(node_map[double_neg_five], five_node)
@@ -76,11 +69,10 @@ class TestRewriteMap(unittest.TestCase):
     double_neg = -(-combined)           # neg(neg(x + y))
     final_expr = zero_node + double_neg  # 0 + (x + y)
 
-    rewritten_sink, node_map = graph_rewrite_map(final_expr, symbolic)
+    node_map = graph_rewrite_map(final_expr, symbolic)
 
     # The final root should be (x_var + y_var).
     expected = x_var + y_var
-    self.assertEqual(rewritten_sink, expected)
 
     # Each sub-expression has its own "final" result.
     # (y + 0) -> y_var
@@ -134,11 +126,7 @@ class TestRewriteMap(unittest.TestCase):
     double_neg_x = -(-x_plus_yz)        # neg(neg(x+(y+z))) -> x+(y+z)
     final_expr = double_neg_x * one_node  # (x+(y+z)) * 1 -> x+(y+z)
 
-    rewritten_sink, node_map = graph_rewrite_map(final_expr, symbolic)
-
-    # Final root should be x + (y + z)
-    expected = x_var + (y_var + z_var)
-    self.assertEqual(rewritten_sink, expected)
+    node_map = graph_rewrite_map(final_expr, symbolic)
 
     # (y + z) is unchanged
     self.assertEqual(node_map[yz_sum], yz_sum)

tinygrad/ops.py

@@ -916,12 +916,13 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False) -> UOp
   _handle_viz(sink, bottom_up)
   return RewriteContext(pm, ctx).bottom_up_rewrite(sink) if bottom_up else RewriteContext(pm, ctx).rewrite(sink)
 
-def graph_rewrite_map(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False) -> tuple[UOp, dict[UOp, UOp]]:
+def graph_rewrite_map(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False) -> dict[UOp, UOp]:
   _handle_viz(sink, bottom_up)
   rewrite_ctx = RewriteContext(pm, ctx)
   rewritten_sink = rewrite_ctx.bottom_up_rewrite(sink) if bottom_up else rewrite_ctx.rewrite(sink)
+  assert rewrite_ctx.replace[sink] == rewritten_sink
   # TODO: is the replace dict correct?
-  return rewritten_sink, rewrite_ctx.replace
+  return rewrite_ctx.replace
 
 # ***** uop type spec *****
 

tinygrad/tensor.py

@@ -225,7 +225,7 @@ class Tensor(SimpleMathTrait):
     scheduled_uops = flatten([x.lazydata.lbs for x in (self,)+lst])
     schedule, var_vals = create_schedule_with_vars(scheduled_uops)
     sink = UOp.sink(*scheduled_uops)
-    _, uop_map = graph_rewrite_map(sink, _substitute, becomes_map, bottom_up=True)
+    uop_map = graph_rewrite_map(sink, _substitute, becomes_map, bottom_up=True)
     for k,v in uop_map.items():
       if (tt:=tensor_map.get(k)) is not None:
         tt.lazydata = v