Merge pull request #176 from IntelPython/update_randint_untyped

Update `randint_untyped` to reuse `randint` implementation

Author: ndgrigorian

CHANGELOG.md

@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Sped up `randint` for `bool`, `uint8`, `int8`, `uint16` and `int16`; generated values are unchanged [gh-172](https://github.com/IntelPython/mkl_random/pull/172)
 * Raised the minimum build-time `Cython` requirement to `3.1.0`, the first release providing the `freethreading_compatible` directive [gh-159](https://github.com/IntelPython/mkl_random/pull/159)
 * Extended the `memcpy`-based fast path of `shuffle` to multi-dimensional `ndarray` inputs whose first-axis items are contiguous, which is also much faster than the previous buffered path [gh-159](https://github.com/IntelPython/mkl_random/pull/159)
+* Updated the `randint_untyped` implementation, which now supports `array_like` (broadcastable) `low`/`high` bounds [gh-176](https://github.com/IntelPython/mkl_random/pull/176)
 
 ### Fixed
 * Fixed `uniform` to return a Python `float` for scalar bounds with `size=None` instead of a 0-d array [gh-167](https://github.com/IntelPython/mkl_random/pull/167)
@@ -28,6 +29,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Fixed a memory leak in `logseries`, which leaked a temporary buffer on every call [gh-159](https://github.com/IntelPython/mkl_random/pull/159)
 * Fixed `randint` with scalar bounds above `INT_MAX`, `tomaxint` and `bytes` returning uninitialized data for `WH`, `MCG31`, `R250` and `MRG32K3A`, which lack `viRngUniformBits` support [gh-175](https://github.com/IntelPython/mkl_random/pull/175)
 * Fixed 64-bit integer generation returning zeros or crashing with `PHILOX4X32X10` and `ARS5` for requests of `2**30` elements or more [gh-175](https://github.com/IntelPython/mkl_random/pull/175)
+* Fixed `randint_untyped` raising `OverflowError` for bounds outside the C `long` range, e.g. `2**40` on Windows [gh-176](https://github.com/IntelPython/mkl_random/pull/176)
 
 ### Removed
 * Removed the `python-gil` constraint from the conda recipes, which pinned `mkl_random` to GIL-enabled Python 3.14 builds [gh-159](https://github.com/IntelPython/mkl_random/pull/159)

mkl_random/mklrand.pyx

@@ -323,12 +323,6 @@ cdef extern from "mkl_distributions.h":
     ) noexcept nogil
 
     # random integers madness
-    void irk_discrete_uniform_vec(
-        irk_state *state, cnp.npy_intp len, int *res, int low, int high
-    ) noexcept nogil
-    void irk_discrete_uniform_long_vec(
-        irk_state *state, cnp.npy_intp len, long *res, long low, long high
-    ) noexcept nogil
     void irk_rand_bool_vec(
         irk_state *state,
         cnp.npy_intp len,
@@ -7066,23 +7060,24 @@ cdef class MKLRandomState(_MKLRandomState):
 
     def randint_untyped(self, low, high=None, size=None):
         """
-        randint_untyped(low, high=None, size=None, dtype=int)
+        randint_untyped(low, high=None, size=None)
 
         Return random integers from `low` (inclusive) to `high` (exclusive).
 
-        Return random integers from the "discrete uniform" distribution of
-        the specified dtype in the "half-open" interval [`low`, `high`). If
-        `high` is None (the default), then results are from [0, `low`).
+        Same as `randint`, except that the result dtype is not selectable:
+        `int32` is used when both bounds fit it, `int64` otherwise.
 
         Parameters
         ----------
-        low : int
+        low : int or array_like of ints
             Lowest (signed) integer to be drawn from the distribution (unless
             ``high=None``, in which case this parameter is the *highest* such
-            integer).
-        high : int, optional
+            integer). If an array is given, it must broadcast with `high` (and
+            with `size`, if provided).
+        high : int or array_like of ints, optional
             If provided, one above the largest (signed) integer to be drawn
             from the distribution (see above for behavior if ``high=None``).
+            If an array is given, it must broadcast with `low`.
         size : int or tuple of ints, optional
             Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
             ``m * n * k`` samples are drawn.  Default is None, in which case a
@@ -7096,79 +7091,35 @@ cdef class MKLRandomState(_MKLRandomState):
 
         See Also
         --------
-        random.random_integers : similar to `randint`, only for the closed
-            interval [`low`, `high`], and 1 is the lowest value if `high` is
-            omitted. In particular, this other one is the one to use to generate
-            uniformly distributed discrete non-integers.
+        randint : same distribution, with a selectable result dtype.
 
         Examples
         --------
-        >>> mkl_random.randint(2, size=10)
-        array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0])
-        >>> mkl_random.randint(1, size=10)
-        array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
-
-        Generate a 2 x 4 array of ints between 0 and 4, inclusive:
-
-        >>> mkl_random.randint(5, size=(2, 4))
-        array([[4, 0, 2, 1],
-               [3, 2, 2, 0]])
+        >>> mkl_random.RandomState().randint_untyped(5, size=(2, 4))
+        array([[4, 0, 2, 1], # random
+               [3, 2, 2, 0]], dtype=int32)
 
         """
-        cdef long lo, hi
-        cdef long *array_long_data
-        cdef int * array_int_data
-        cdef cnp.ndarray array "arrayObject"
-        cdef cnp.npy_intp length
-        cdef int rv_int
-        cdef long rv_long
-
         if high is None:
-            lo = 0
-            hi = low
+            high = low
+            low = 0
+
+        # untyped: narrowest of int32/int64 holding both bounds,
+        # `initial` guards empty ones
+        lo_min = low if np.isscalar(low) else np.min(low, initial=0)
+        hi_max = high if np.isscalar(high) else np.max(high, initial=0)
+        if (-2**31 <= lo_min) and (hi_max <= 2**31 - 1):
+            _dtype = np.int32
         else:
-            lo = low
-            hi = high
+            _dtype = np.int64
 
-        if lo >= hi :
-            raise ValueError("low >= high")
+        res = self.randint(low, high, size=size, dtype=_dtype)
 
-        if ((<int> lo) == lo) and ((<int>hi) == hi):
-            if size is None:
-                with self.lock, nogil:
-                    irk_discrete_uniform_vec(
-                        self.internal_state, 1, &rv_int, <int>lo, <int>hi
-                    )
-                return rv_int
-            else:
-                array = <cnp.ndarray>np.empty(size, np.int32)
-                length = cnp.PyArray_SIZE(array)
-                array_int_data = <int*>cnp.PyArray_DATA(array)
-                with self.lock, nogil:
-                    irk_discrete_uniform_vec(
-                        self.internal_state,
-                        length,
-                        array_int_data,
-                        <int>lo,
-                        <int>hi
-                    )
-                return array
-        else:
-            if size is None:
-                with self.lock, nogil:
-                    irk_discrete_uniform_long_vec(
-                        self.internal_state, 1, &rv_long, lo, hi
-                    )
-                return rv_long
-            else:
-                array = <cnp.ndarray>np.empty(size, int)
-                length = cnp.PyArray_SIZE(array)
-                array_long_data = <long*>cnp.PyArray_DATA(array)
-                with self.lock, nogil:
-                    irk_discrete_uniform_long_vec(
-                        self.internal_state, length, array_long_data, lo, hi
-                    )
-                return array
+        # a single sample has always been a Python int
+        if size is None and res.ndim == 0:
+            return int(res)
+
+        return res
 
     def multinormal_cholesky(self, mean, ch, size=None, method=ICDF):
         """

mkl_random/src/mkl_distributions.cpp

@@ -1594,32 +1594,6 @@ void irk_logseries_vec(irk_state *state,
     mkl_free(Uvec);
 }
 
-/* samples discrete uniforms from [low, high) */
-void irk_discrete_uniform_vec(irk_state *state,
-                              npy_intp len,
-                              int *res,
-                              const int low,
-                              const int high)
-{
-    int err = 0;
-
-    if (len < 1)
-        return;
-
-    while (len > MKL_INT_MAX) {
-        err = viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, state->stream,
-                           MKL_INT_MAX, res, low, high);
-        assert(err == VSL_STATUS_OK);
-
-        res += MKL_INT_MAX;
-        len -= MKL_INT_MAX;
-    }
-
-    err = viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, state->stream, len, res, low,
-                       high);
-    assert(err == VSL_STATUS_OK);
-}
-
 /*
  * Bulk source of raw uniform words for the bounded-integer
  * routines below, overloaded on the word type (32/64-bit). BRNGs that lack
@@ -1717,85 +1691,6 @@ void irk_uniform_bits32_vec(irk_state *state, npy_intp len, npy_uint32 *res)
     irk_uniform_bits_vec(state, len, res);
 }
 
-void irk_discrete_uniform_long_vec(irk_state *state,
-                                   npy_intp len,
-                                   long *res,
-                                   const long low,
-                                   const long high)
-{
-    int err = 0;
-    unsigned long max;
-    npy_intp i = 0;
-
-    if (len < 1)
-        return;
-
-    while (len > MKL_INT_MAX) {
-        irk_discrete_uniform_long_vec(state, MKL_INT_MAX, res, low, high);
-
-        res += MKL_INT_MAX;
-        len -= MKL_INT_MAX;
-    }
-
-    max = ((unsigned long)high) - ((unsigned long)low) - 1UL;
-    if (max == 0) {
-        DIST_PRAGMA_VECTOR
-        for (i = 0; i < len; ++i)
-            res[i] = low;
-
-        return;
-    }
-
-    if (max <= (unsigned long)INT_MAX) {
-        int *buf = (int *)mkl_malloc(len * sizeof(int), 64);
-        assert(buf != nullptr);
-
-        err = viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, state->stream, len, buf,
-                           -1, (int)max);
-        assert(err == VSL_STATUS_OK);
-
-        DIST_PRAGMA_VECTOR
-        for (i = 0; i < len; ++i)
-            res[i] = low + ((long)buf[i]) + 1L;
-
-        mkl_free(buf);
-    }
-    else {
-        unsigned long mask = max;
-        unsigned long *buf = nullptr;
-        int n_accepted;
-
-        /* Smallest bit mask >= max */
-        mask |= mask >> 1;
-        mask |= mask >> 2;
-        mask |= mask >> 4;
-        mask |= mask >> 8;
-        mask |= mask >> 16;
-#if ULONG_MAX > 0xffffffffUL
-        mask |= mask >> 32;
-#endif
-
-        buf = (unsigned long *)mkl_malloc(len * sizeof(long), 64);
-        assert(buf != nullptr);
-        n_accepted = 0;
-
-        while (n_accepted < len) {
-            int k, batchSize = len - n_accepted;
-
-            irk_uniform_bits_vec(state, batchSize, (npy_uint64 *)buf);
-
-            for (k = 0; k < batchSize; ++k) {
-                unsigned long value = buf[k] & mask;
-                if (value <= max) {
-                    res[n_accepted++] = low + value;
-                }
-            }
-        }
-
-        mkl_free(buf);
-    }
-}
-
 void irk_ulong_vec(irk_state *state, npy_intp len, unsigned long *res)
 {
     if (len < 1)

mkl_random/src/mkl_distributions.h

@@ -249,18 +249,6 @@ extern "C"
     extern void
         irk_uniform_bits32_vec(irk_state *state, npy_intp len, npy_uint32 *res);
 
-    extern void irk_discrete_uniform_vec(irk_state *state,
-                                         npy_intp len,
-                                         int *res,
-                                         const int low,
-                                         const int high);
-
-    extern void irk_discrete_uniform_long_vec(irk_state *state,
-                                              npy_intp len,
-                                              long *res,
-                                              const long low,
-                                              const long high);
-
     extern void irk_rand_int64_vec(irk_state *state,
                                    npy_intp len,
                                    npy_int64 *res,

mkl_random/tests/test_random.py

@@ -498,6 +498,36 @@ def test_array_bounds_errors(self):
         assert_raises(ValueError, rnd.randint, [0, 0], [5, 6], (4,))
         assert_raises(ValueError, rnd.randint, [3, 4], [9, 10], ())
 
+    def test_untyped_result_type(self):
+        rs = rnd.MKLRandomState(0)
+        assert rs.randint_untyped(0, 100, size=10).dtype == np.int32
+        assert rs.randint_untyped(0, 2**31, size=10).dtype == np.int64
+        assert rs.randint_untyped(-(2**31) - 1, 0, size=10).dtype == np.int64
+        assert type(rs.randint_untyped(5)) is int
+        assert rs.randint_untyped(5, size=()).shape == ()
+
+    def test_untyped_array_bounds(self):
+        low, high = [0, 10, 20], [10, 20, 30]
+        vals = rnd.MKLRandomState(0).randint_untyped(low, high)
+        assert vals.dtype == np.int32
+        assert np.all(vals >= low)
+        assert np.all(vals < high)
+
+    def test_untyped_errors(self):
+        rs = rnd.MKLRandomState(0)
+        assert_raises(ValueError, rs.randint_untyped, 5, 5)
+        assert_raises(ValueError, rs.randint_untyped, [0, 5], [5, 5])
+
+    @pytest.mark.parametrize("brng", _ALL_BRNGS)
+    def test_untyped_range_above_int_max(self, brng):
+        # a range wider than INT_MAX used to take a separate C long path
+        hi = 2**40
+        y = rnd.MKLRandomState(0, brng=brng).randint_untyped(0, hi, size=100000)
+        assert y.dtype == np.int64
+        assert y.min() >= 0
+        assert int(y.max()) < hi
+        assert len(np.unique(y)) > 99000
+
 
 class RandomDistData(NamedTuple):
     seed: int