randint accept ints for 'size' (#916)

* fix size bug in randint

* use else

* return scalar size==None

Co-authored-by: Claudia Comito <[email protected]>

CHANGELOG.md

@@ -10,6 +10,7 @@
 - [#846](https://github.com/helmholtz-analytics/heat/pull/846) Fixed an issue in `_reduce_op` when axis and keepdim were set.
 - [#846](https://github.com/helmholtz-analytics/heat/pull/846) Fixed an issue in `min`, `max` where DNDarrays with empty processes can't be computed.
 - [#868](https://github.com/helmholtz-analytics/heat/pull/868) Fixed an issue in `__binary_op` where data was falsely distributed if a DNDarray has single element.
+- [#916](https://github.com/helmholtz-analytics/heat/pull/916) Fixed an issue in `random.randint` where the size parameter does not accept ints.
 
 ## Feature Additions
 

heat/core/random.py

@@ -527,10 +527,14 @@ def randint(
 
     # sanitize shape
     if size is None:
-        size = (1,)
-    shape = tuple(int(ele) for ele in size)
-    if not all(ele > 0 for ele in shape):
-        raise ValueError("negative dimensions are not allowed")
+        size = ()
+    try:
+        shape = tuple(int(ele) for ele in size)
+    except TypeError:
+        shape = (int(size),)
+    else:
+        if not all(ele >= 0 for ele in shape):
+            raise ValueError("negative dimensions are not allowed")
 
     # sanitize the data type
     if dtype is None:

heat/core/tests/test_random.py

@@ -240,6 +240,14 @@ def test_randint(self):
         b = ht.zeros((10,), dtype=ht.int64, split=0)
         self.assertTrue(ht.equal(a, b))
 
+        # size parameter allows int arguments
+        a = ht.random.randint(1, size=10, split=0, dtype=ht.int64)
+        self.assertTrue(ht.equal(a, b))
+
+        # size is None
+        a = ht.random.randint(0, 10)
+        self.assertEqual(a.shape, ())
+
         # Two arrays with the same seed and same number of elements have the same random values
         ht.random.seed(13579)
         shape = (15, 13, 9, 21, 65)