diff --git a/test/test_nn.py b/test/test_nn.py index 9619ce04b1..736d9e3e1d 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -41,9 +41,9 @@ class TestNN(unittest.TestCase): toutt = tbn(torch.tensor(inn.cpu().numpy())) # close - np.testing.assert_allclose(outt.numpy(), toutt.detach().numpy(), rtol=5e-4) + np.testing.assert_allclose(outt.numpy(), toutt.detach().numpy(), rtol=5e-4, atol=1e-6) - np.testing.assert_allclose(bn.running_mean.numpy(), tbn.running_mean.detach().numpy(), rtol=1e-5) + np.testing.assert_allclose(bn.running_mean.numpy(), tbn.running_mean.detach().numpy(), rtol=1e-5, atol=1e-6) # TODO: this is failing # np.testing.assert_allclose(bn.running_var.numpy(), tbn.running_var.detach().numpy(), rtol=1e-5) diff --git a/test/unit/test_shapetracker.py b/test/unit/test_shapetracker.py index a5802bcc47..a96f5d11de 100644 --- a/test/unit/test_shapetracker.py +++ b/test/unit/test_shapetracker.py @@ -2,7 +2,7 @@ import unittest import numpy as np from tinygrad.helpers import prod -from tinygrad.shape import ShapeTracker, View, ZeroView +from tinygrad.shape import ShapeTracker, View, ZeroView, merge_views from tinygrad.codegen.gpu import to_image_idx def shapetracker_getitem(st, val): @@ -75,6 +75,12 @@ class TestImageShapeTracker(unittest.TestCase): offsets = [0,32,64] """ + new_view = merge_views( + View((1, 66, 130, 32, 1, 1), (0, 4096, 32, 1, 0, 0), -4128), + View((64, 32, 8, 3, 3), (4160, 128, 4, 4160, 32), 0) + ) + print(new_view) + st = ShapeTracker(shape=(64, 32, 8, 3, 3), views=[ View((1, 66, 130, 32, 1, 1), (0, 4096, 32, 1, 0, 0), -4128), ZeroView((1, 64, 128, 32, 1, 1), ((0, 1), (-1, 65), (-1, 129), (0, 32), (0, 1), (0, 1))),