BUG: random: randint returns integers, not floats

Author: ev-br

torch_np/random.py

@@ -15,7 +15,8 @@
 
 _default_dtype = _default_float_type.torch_dtype
 
-__all__ = ["seed", "random_sample", "sample", "random", "rand", "randn", "normal"]
+__all__ = ["seed", "random_sample", "sample", "random", "rand", "randn",
+           "normal", "choice", "randint", "shuffle", "uniform"]
 
 
 def array_or_scalar(values, py_type=float):
@@ -79,7 +80,7 @@ def randint(low, high=None, size=None):
     if high is None:
         low, high = 0, low
     values = torch.randint(low, high, size=size)
-    return array_or_scalar(values)
+    return array_or_scalar(values, int)
 
 
 def choice(a, size=None, replace=True, p=None):