fix shapetracker test

docs/abstractions.py

@@ -384,5 +384,4 @@ print(expr.min, expr.max)  # 0 20
 # this is just "(a*2)"
 # since b only has a range from 0-10, it can't affect the output
 
-
 # %%

test/test_jit.py

@@ -9,7 +9,7 @@ class TestJit(unittest.TestCase):
   def test_simple_jit(self):
     @TinyJit
     def add(a, b): return (a+b).realize()
-    for _ in range(3):
+    for _ in range(5):
       a = Tensor.randn(10, 10)
       b = Tensor.randn(10, 10)
       c = add(a, b)
@@ -18,7 +18,7 @@ class TestJit(unittest.TestCase):
   def test_kwargs_jit(self):
     @TinyJit
     def add_kwargs(first, second): return (first+second).realize()
-    for _ in range(3):
+    for _ in range(5):
       a = Tensor.randn(10, 10)
       b = Tensor.randn(10, 10)
       c = add_kwargs(first=a, second=b)
@@ -27,12 +27,12 @@ class TestJit(unittest.TestCase):
   def test_array_jit(self):
     @TinyJit
     def add_array(a, arr): return (a+arr[0]).realize()
-    for i in range(3):
+    for i in range(5):
       a = Tensor.randn(10, 10)
       b = Tensor.randn(10, 10)
       a.realize(), b.realize()
       c = add_array(a, [b])
-      if i == 2:
+      if i >= 2:
         # should fail once jitted since jit can't handle arrays
         np.testing.assert_equal(np.any(np.not_equal(c.numpy(),a.numpy()+b.numpy())), True)
       else:

test/unit/test_shapetracker.py

@@ -23,31 +23,31 @@ class CheckingShapeTracker:
   def simplify(self): self.st.simplify()
 
   def reshape(self, new_shape):
-    self.st._reshape(new_shape)
+    self.st.reshape(new_shape)
     self.t = self.t.reshape(new_shape)
 
   def permute(self, axis):
-    self.st._permute(axis)
+    self.st.permute(axis)
     self.t = np.transpose(self.t, axis)
 
   def expand(self, new_shape):
-    self.st._expand(new_shape)
+    self.st.expand(new_shape)
     self.t = np.broadcast_to(self.t, new_shape)
 
   def flip(self, axis):
-    self.st._stride(tuple(-1 if i in axis else 1 for i in range(len(self.shape))))
+    self.st.stride(tuple(-1 if i in axis else 1 for i in range(len(self.shape))))
     self.t = np.flip(self.t, axis)
 
   def shrink(self, arg):
-    self.st._shrink(arg)
+    self.st.shrink(arg)
     self.t = self.t[tuple([slice(x[0], x[1]) for x in arg])]
 
   def pad(self, arg):
-    self.st._pad(arg)
+    self.st.pad(arg)
     self.t = np.pad(self.t, arg, constant_values=-1)
 
   def stride(self, arg):
-    self.st._stride(arg)
+    self.st.stride(arg)
     self.t = self.t[tuple([slice(None, None, x) for x in arg])]
 
   def __getitem__(self, val):