From 25be7f745de71b9ac45fdcc496bd032ecdd76702 Mon Sep 17 00:00:00 2001
From: Jordan Wright <33017513+jaw566@users.noreply.github.com>
Date: Sat, 26 Aug 2023 01:59:53 -0400
Subject: [PATCH] Tensor.uniform with dtype=int bug fix (#1593)

---
 test/test_randomness.py | 9 +++++++++
 tinygrad/tensor.py      | 4 +++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/test/test_randomness.py b/test/test_randomness.py
index 37300c54..53be0109 100644
--- a/test/test_randomness.py
+++ b/test/test_randomness.py
@@ -5,6 +5,7 @@ import torch
 from tinygrad.tensor import Tensor
 import tinygrad.nn as nn
 import pytest
+from tinygrad.helpers import dtypes
 
 pytestmark = pytest.mark.webgpu
 
@@ -49,6 +50,13 @@ def normal_test(func, shape=(20, 23), alpha=0.05):
   y = np.random.randn(*shape).flatten()
   return kstest(x, y) >= alpha
 
+def equal_distrib_ints(tiny_func, numpy_func, shape=(20, 23), low=-100, high=100, dtype=dtypes.int32, alpha=0.05):
+  Tensor.manual_seed(1337)
+  np.random.seed(1337)
+  x = tiny_func(*shape, low=low, high=high, dtype=dtype).cpu().numpy().flatten()
+  y = numpy_func(shape).flatten()
+  return kstest(x, y) >= alpha
+
 def equal_distribution(tiny_func, torch_func, numpy_func=None, shape=(20, 23), alpha=0.05):
   Tensor.manual_seed(1337)
   torch.manual_seed(1337)
@@ -74,6 +82,7 @@ class TestRandomness(unittest.TestCase):
   def test_uniform(self):
     self.assertFalse(normal_test(Tensor.uniform))
     self.assertTrue(equal_distribution(Tensor.uniform, lambda x: torch.nn.init.uniform_(torch.empty(x), a=-1, b=1), lambda x: np.random.uniform(low=-1, high=1, size=x)))
+    self.assertTrue(equal_distrib_ints(Tensor.uniform, lambda x: np.random.randint(low=-100, high=100, size=x)))
 
   def test_scaled_uniform(self):
     self.assertFalse(normal_test(Tensor.scaled_uniform))
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index 917f14f9..1e4087fc 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -179,7 +179,9 @@ class Tensor:
   def normal(*shape, mean=0.0, std=1.0, **kwargs) -> Tensor: return (std * Tensor.randn(*shape, **kwargs)) + mean
 
   @staticmethod
-  def uniform(*shape, low=-1.0, high=1.0, **kwargs) -> Tensor: return ((high-low) * Tensor.rand(*shape, **kwargs)) + low
+  def uniform(*shape, low=-1.0, high=1.0, **kwargs) -> Tensor:
+    dtype = kwargs.pop("dtype", Tensor.default_type)
+    return ((high-low) * Tensor.rand(*shape, **kwargs)).cast(dtype) + low
 
   @staticmethod
   def scaled_uniform(*shape, **kwargs) -> Tensor: return Tensor.uniform(*shape, **kwargs).mul(math.prod(shape)**-0.5)