more tests pass

test/test_const_folding.py

@@ -1,14 +1,13 @@
 import unittest, math
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.ops import Ops
-from tinygrad.engine.schedule import create_schedule
 from tinygrad.helpers import CI
 import numpy as np
 from tinygrad.device import is_dtype_supported
 
 def _check_ast_count(desired_count:int, t:Tensor):
   # NOTE: this has side effect because everything can be scheduled only once
-  schedule = create_schedule(t.lazydata.lbs)
+  schedule = t.schedule()
   asts = [s for s in schedule if s.ast.op is Ops.SINK]
   assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
 

test/test_nn.py

@@ -8,7 +8,6 @@ from tinygrad.helpers import CI, Context
 from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding
 from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell
 from tinygrad.nn.state import load_state_dict
-from tinygrad.engine.schedule import create_schedule
 from tinygrad.engine.realize import run_schedule
 from tinygrad.device import is_dtype_supported
 
@@ -517,7 +516,7 @@ class TestNN(unittest.TestCase):
     a = Tensor([[1, 5, 9, 11],
                 [12, 19, 8, 1]])
     result = layer(a)
-    schedule = create_schedule([result.lazydata])
+    schedule = result.schedule()
     self.assertEqual(3, len([item for item in schedule if item.ast.op is Ops.SINK]), "first run realizes arange, weight, and embedding")
     run_schedule(schedule)
 
@@ -525,7 +524,7 @@ class TestNN(unittest.TestCase):
                 [4, 5, 6],
                 [7, 8, 9]])
     result = layer(b)
-    schedule = create_schedule([result.lazydata])
+    schedule = result.schedule()
     self.assertEqual(1, len([item for item in schedule if item.ast.op is Ops.SINK]), "second run realizes embedding only")
     run_schedule(schedule)