diff --git a/test/test_randomness.py b/test/test_randomness.py index 99e00ba984..ce79e85f45 100644 --- a/test/test_randomness.py +++ b/test/test_randomness.py @@ -58,6 +58,15 @@ class TestRandomness(unittest.TestCase): self.assertFalse(normal_test(Tensor.rand)) self.assertTrue(equal_distribution(Tensor.rand, torch.rand, lambda x: np.random.rand(*x))) + def test_rand_half(self): + N = 128 + x = Tensor.rand((2, N, N), dtype=dtypes.half).realize().numpy() + ones = np.take(x, np.where(x == 1)) + zeros = np.take(x, np.where(x == 0)) + self.assertTrue(ones.size == 0) + self.assertTrue(zeros.size > 0) + equal_distribution(lambda *x: Tensor.rand(*x, dtype=dtypes.float16), torch.rand, lambda x: np.random.rand(*x), shape=(2, N, N)) + def test_randn(self): self.assertTrue(normal_test(Tensor.randn)) self.assertTrue(equal_distribution(Tensor.randn, torch.randn, lambda x: np.random.randn(*x))) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 891e13075c..1aeccb897c 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1014,5 +1014,6 @@ def custom_random(out:Buffer): Tensor._seed += 1 if DEBUG >= 2: print(f"*** {out.device} rand seed {Tensor._seed} size {out.size:<15d} dtype {out.dtype}") rng = np.random.default_rng(Tensor._seed) - rng_np_buffer = rng.random(size=out.size, dtype=np.float32).astype(dtype=out.dtype.np, copy=False) + if out.dtype == dtypes.half: rng_np_buffer = (rng.integers(low=0, high=2047, size=out.size) / 2048).astype(np.half, copy=False) + else: rng_np_buffer = rng.random(size=out.size, dtype=np.float32).astype(dtype=out.dtype.np, copy=False) out.copyin(rng_np_buffer.data)