diff --git a/test/backend/test_randomness.py b/test/backend/test_randomness.py index 5ac503b812..adb03f0873 100644 --- a/test/backend/test_randomness.py +++ b/test/backend/test_randomness.py @@ -368,8 +368,8 @@ class TestRandomness(unittest.TestCase): @TinyJit def sample_one(): return Tensor(w).multinomial(1, replacement=False).realize() - tiny_samples = [sample_one().item() for _ in range(1000)] - torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(1000)] + tiny_samples = [sample_one().item() for _ in range(400)] + torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(400)] self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples))) w = list(range(32)) @@ -384,8 +384,8 @@ class TestRandomness(unittest.TestCase): @TinyJit def sample_three(): return Tensor(w).multinomial(3, replacement=False).realize() - tiny_draws = np.array([sample_three().numpy() for _ in range(1000)]) - torch_draws = np.array([torch.tensor(w).multinomial(3, replacement=False).numpy() for _ in range(1000)]) + tiny_draws = np.array([sample_three().numpy() for _ in range(400)]) + torch_draws = np.array([torch.tensor(w).multinomial(3, replacement=False).numpy() for _ in range(400)]) for pos in range(3): self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_draws[:, pos]), lambda _: torch.tensor(torch_draws[:, pos]))) @@ -415,7 +415,7 @@ class TestRandomness(unittest.TestCase): def test_rand_chain(self): # NOTE: this fails if property propagates deeper than stack limit for _ in range(833): Tensor.rand(1) - Tensor.rand(1).realize() + Tensor.rand(1).schedule_linear() def test_random_counter_overflow(self): device = Device.DEFAULT diff --git a/test/backend/test_schedule.py b/test/backend/test_schedule.py index 102c772565..cf327cc06c 100644 --- a/test/backend/test_schedule.py +++ b/test/backend/test_schedule.py @@ -5,7 +5,7 @@ import gc, unittest, functools import numpy as np from typing import cast -from hypothesis import assume, given, settings, strategies as strat +from hypothesis import assume, given, strategies as strat from tinygrad import nn, dtypes, Device, Tensor, Variable from tinygrad.device import is_dtype_supported @@ -113,12 +113,6 @@ class TestSchedule(unittest.TestCase): run_linear(*check_schedule(b, 1)) np.testing.assert_allclose(b.numpy(), np.broadcast_to(a.numpy().astype(np.float16), (2, 4, 4))+2, rtol=1e-3) - def test_indexing_scalars_simple(self): - X = Tensor.randn(2, 2).realize() - xt = X[Tensor(1)][Tensor(0)] - run_linear(*check_schedule(xt, 1)) - np.testing.assert_equal(xt.numpy(), X.numpy()[1][0]) - @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI") def test_add_chain_buffers(self): N = 31 @@ -130,14 +124,14 @@ class TestSchedule(unittest.TestCase): root = root + functools.reduce(lambda a,b:a+b, bufs[i:i+X]) self.assertEqual(root.item(), sum(range(N))) - @given(strat.sampled_from(range(2,4)), strat.sampled_from(range(2,4)), strat.sampled_from(range(0,4)), strat.sampled_from(range(0,4))) - @settings(deadline=None) - def test_indexing_scalars(self, x, y, a, b): - assume(a