test tiny passes

test/unit/test_rewrite_map.py

@@ -1,20 +1,21 @@
 import unittest
 from tinygrad import dtypes, Tensor
 from tinygrad.ops import UOp, symbolic, graph_rewrite_map, _substitute
+from test.unit.test_tensor_uop_representation import is_pattern, realized_pattern
 
 class TestTensorMutates(unittest.TestCase):
   def test_mutate_add(self):
     a = Tensor([1,2,3])
     b = Tensor([4,5,6])
     ret = a+b
+    pa = a.lazydata
+    pb = b.lazydata
+    pr = ret.lazydata
     ret.schedule()
-    print(a.lazydata)
-    print(b.lazydata)
-    print(ret.lazydata)
-    from tinygrad.ops import becomes_map
-    print(becomes_map)
-    for k,v in becomes_map.items():
-      print(k,v)
+    self.assertIsNot(pa, a.lazydata)
+    self.assertIsNot(pb, b.lazydata)
+    self.assertIsNot(pr, ret.lazydata)
+    for t in [a,b,ret]: is_pattern(t, realized_pattern)
 
 class TestRewriteMap(unittest.TestCase):
   def test_substitute(self):

tinygrad/tensor.py

@@ -14,6 +14,13 @@ from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.memory import memory_planner
 from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
 
+# *** Tensors are containers for UOps ***
+
+tensor_map: weakref.WeakValueDictionary[UOp, Tensor] = weakref.WeakValueDictionary()
+def update_tensor_map(t:Tensor):
+  # TODO: multi
+  tensor_map[t.lazydata] = t
+
 # **** start with two base classes, Tensor and Function ****
 
 class Function:
@@ -33,6 +40,7 @@ class Function:
     ret = Tensor.__new__(Tensor)
     ret.lazydata, ret.requires_grad, ret.grad = ctx.forward(*[t.lazydata for t in x], **kwargs), ctx.requires_grad, None
     ret._ctx = ctx if ctx.requires_grad and not Tensor.no_grad else None  # used by autograd engine
+    update_tensor_map(ret)
     return ret
 
 import tinygrad.function as F
@@ -105,11 +113,6 @@ def _masked_setitem(target:Tensor, values:Tensor, mask:Tensor, axes:tuple[int, .
 
 ReductionStr = Literal["mean", "sum", "none"]
 
-tensor_map: weakref.WeakValueDictionary[UOp, Tensor] = weakref.WeakValueDictionary()
-def update_tensor_map(t:Tensor):
-  # TODO: multi
-  tensor_map[t.lazydata] = t
-
 class Tensor(SimpleMathTrait):
   """
   A `Tensor` is a multi-dimensional matrix containing elements of a single data type.
@@ -224,10 +227,11 @@ class Tensor(SimpleMathTrait):
     """
     scheduled_uops = flatten([x.lazydata.lbs for x in (self,)+lst])
     schedule, var_vals = create_schedule_with_vars(scheduled_uops)
+    rewrite_map = graph_rewrite_map(UOp.sink(*scheduled_uops), _substitute, becomes_map, bottom_up=True)
 
     # apply becomes_map
     # TODO: add children to scheduled_uops
-    for k,v in graph_rewrite_map(UOp.sink(*scheduled_uops), _substitute, becomes_map).items():
+    for k,v in rewrite_map.items():
       if (tt:=tensor_map.get(k)) is not None:
         tt.lazydata = v
         tensor_map[v] = tt