small changes from block linearizer [pr] (#7888)

* small changes from block linearizer [pr]

* fix test_gc

test/external/external_benchmark_schedule.py

@@ -48,7 +48,8 @@ if __name__ == "__main__":
                 rewritten_uops.append(full_graph_rewrite(u, k.opts))
             uops = rewritten_uops
         if getenv("LINEARIZE", 1):
-          with Timing("***** model linearize in "): uops = [linearize_uop(u) for u in uops]
+          with Profiling(PROFILE >= 2):
+            with Timing("***** model linearize in "): uops = [linearize_uop(u) for u in uops]
           print(sum(len(u) for u in uops))
           if getenv("SRC", 0):
             renderer = Device[Device.DEFAULT].renderer

test/test_gc.py

@@ -16,30 +16,32 @@ class TestGC(unittest.TestCase):
 
   def test_gc(self):
     Tensor.manual_seed(0)
+    base = tensors_allocated()
     a = Tensor.rand(4, 4, requires_grad=True)
     b = Tensor.zeros(4, 4, requires_grad=True)
     (a*b).mean().backward()
-    assert (tensors_allocated() > 0)
+    assert (tensors_allocated()-base > 0)
     del a,b
-    assert (tensors_allocated() == 2) # one for Tensor._device_rng_counters, and one for Tensor._device_seeds
+    assert (tensors_allocated()-base == 2) # one for Tensor._device_rng_counters, and one for Tensor._device_seeds
     Tensor.manual_seed(0)
 
   def test_gc_complex(self):
     Tensor.manual_seed(0)
+    base = tensors_allocated()
     a = Tensor(np.zeros((4, 4), dtype=np.float32), requires_grad=True)
     b = Tensor.rand(4, 4, requires_grad=True)
-    assert (tensors_allocated() == 5)
+    assert (tensors_allocated()-base == 5)
     (a*b).mean().backward()
-    assert (tensors_allocated() == 6)
+    assert (tensors_allocated()-base == 6)
     del b
-    assert (tensors_allocated() == 4)
+    assert (tensors_allocated()-base == 4)
     b = Tensor(np.zeros((4, 4), dtype=np.float32), requires_grad=True)
     print(tensors_allocated())
     (a*b).mean().backward()
     print(tensors_allocated())
-    assert (tensors_allocated() == 6)
+    assert (tensors_allocated()-base == 6)
     del b
-    assert (tensors_allocated() == 4)
+    assert (tensors_allocated()-base == 4)
     Tensor.manual_seed(0)
 
   def test_schedule_gc(self):

tinygrad/helpers.py

@@ -66,7 +66,7 @@ def get_child(obj, key):
     elif isinstance(obj, dict): obj = obj[k]
     else: obj = getattr(obj, k)
   return obj
-def word_wrap(x, wrap=80): return x if len(x) <= wrap else (x[0:wrap] + "\n" + word_wrap(x[wrap:], wrap))
+def word_wrap(x, wrap=80): return x if len(x) <= wrap or '\n' in x[0:wrap] else (x[0:wrap] + "\n" + word_wrap(x[wrap:], wrap))
 
 # for length N coefficients `p`, returns p[0] * x**(N-1) + p[1] * x**(N-2) + ... + p[-2] * x + p[-1]
 def polyN(x:T, p:List[float]) -> T: return functools.reduce(lambda acc,c: acc*x+c, p, 0.0)  # type: ignore

tinygrad/ops.py

@@ -548,7 +548,7 @@ def get_location() -> Tuple[str, int]:
   frm = sys._getframe(1)
   # find the real frame in the file that has the UPat, TODO: is there a better way to do this?
   while frm.f_back is not None and pathlib.Path(frm.f_back.f_code.co_filename).name in {"ops.py", "uopgraph.py", "schedule.py",
-                                                                                        "lowerer.py", "cstyle.py"}:
+                                                                                        "lowerer.py", "cstyle.py", "linearize.py"}:
     frm = frm.f_back
   return frm.f_code.co_filename, frm.f_lineno
 @functools.lru_cache(None)