diff --git a/docs/abstractions.py b/docs/abstractions.py index 2051e5f2e0..117436cf5d 100644 --- a/docs/abstractions.py +++ b/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 - # %% diff --git a/test/test_jit.py b/test/test_jit.py index 61838530e0..a97dd5a589 100644 --- a/test/test_jit.py +++ b/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: diff --git a/test/unit/test_shapetracker.py b/test/unit/test_shapetracker.py index 59c437c513..5129fd1a99 100644 --- a/test/unit/test_shapetracker.py +++ b/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):