From ff953f1c01bf982ffa39c196333210cd80ca880f Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Tue, 1 Sep 2026 04:13:28 -0700 Subject: [PATCH 01/14] Add irk_rand__broadcast for randint array-like bounds --- mkl_random/src/mkl_distributions.cpp | 208 +++++++++++++++++++++++++++ mkl_random/src/mkl_distributions.h | 46 ++++++ 2 files changed, 254 insertions(+) diff --git a/mkl_random/src/mkl_distributions.cpp b/mkl_random/src/mkl_distributions.cpp index 8fa67d96..b8311c28 100644 --- a/mkl_random/src/mkl_distributions.cpp +++ b/mkl_random/src/mkl_distributions.cpp @@ -2151,6 +2151,214 @@ void irk_rand_int64_vec(irk_state *state, res[i] = res[i] + lo; } +/* + * Bulk generators of raw uniform words used by the broadcasted bounded-integer + * routines below. Overloaded on the word type so that the masked-rejection + * template can pick 32- or 64-bit words at compile time. + */ +static inline void + irk_uniform_bits_vec(irk_state *state, npy_intp len, npy_uint32 *buf) +{ + int err = 0; + + while (len > 0) { + MKL_INT c = (len > MKL_INT_MAX) ? (MKL_INT)MKL_INT_MAX : (MKL_INT)len; + err = viRngUniformBits32(VSL_RNG_METHOD_UNIFORMBITS32_STD, + state->stream, c, (unsigned int *)buf); + assert(err == VSL_STATUS_OK); + buf += c; + len -= c; + } +} + +static inline void + irk_uniform_bits_vec(irk_state *state, npy_intp len, npy_uint64 *buf) +{ + int err = 0; + + while (len > 0) { + MKL_INT c = (len > MKL_INT_MAX) ? (MKL_INT)MKL_INT_MAX : (MKL_INT)len; + err = viRngUniformBits64(VSL_RNG_METHOD_UNIFORMBITS64_STD, + state->stream, c, (unsigned MKL_INT64 *)buf); + assert(err == VSL_STATUS_OK); + buf += c; + len -= c; + } +} + +/* Smallest bit mask (2^k - 1) that is >= rng. */ +template +static inline UT irk_gen_mask(UT rng) +{ + UT mask = rng; + unsigned int s = 0; + + for (s = 1; s < sizeof(UT) * 8; s <<= 1) + mask |= mask >> s; + + return mask; +} + +/* + * Draw res[i] uniformly from [low[i], hi[i]] (inclusive) using per-element + * masked rejection, the same algorithm as irk_rand_uint64_vec but with + * per-element bounds. Words are generated in bulk by MKL; rejected elements + * are gathered into `idx` (allocated lazily) and retried on the next round. + * T is the result type, UT its unsigned counterpart, WT the raw-word type. + */ +template +static void irk_rand_bounded_broadcast(irk_state *state, + npy_intp len, + T *res, + const T *low, + const T *hi) +{ + npy_intp i = 0; + npy_intp k = 0; + npy_intp n_pending = 0; + npy_intp *idx = nullptr; + WT *words = nullptr; + + if (len < 1) + return; + + words = (WT *)mkl_malloc(len * sizeof(WT), 64); + assert(words != nullptr); + + irk_uniform_bits_vec(state, len, words); + + for (i = 0; i < len; ++i) { + UT rng = ((UT)hi[i]) - ((UT)low[i]); + UT value = ((UT)words[i]) & irk_gen_mask(rng); + + if (value <= rng) { + res[i] = (T)(((UT)low[i]) + value); + } + else { + if (idx == nullptr) { + idx = (npy_intp *)mkl_malloc(len * sizeof(npy_intp), 64); + assert(idx != nullptr); + } + idx[n_pending++] = i; + } + } + + while (n_pending > 0) { + npy_intp w = 0; + + irk_uniform_bits_vec(state, n_pending, words); + + for (k = 0; k < n_pending; ++k) { + npy_intp j = idx[k]; + UT rng = ((UT)hi[j]) - ((UT)low[j]); + UT value = ((UT)words[k]) & irk_gen_mask(rng); + + if (value <= rng) { + res[j] = (T)(((UT)low[j]) + value); + } + else { + /* keep this element pending; w <= k so idx[k] is read first */ + idx[w++] = j; + } + } + n_pending = w; + } + + if (idx != nullptr) + mkl_free(idx); + mkl_free(words); +} + +void irk_rand_bool_broadcast(irk_state *state, + npy_intp len, + npy_bool *res, + const npy_bool *low, + const npy_bool *hi) +{ + irk_rand_bounded_broadcast(state, len, res, + low, hi); +} + +void irk_rand_int8_broadcast(irk_state *state, + npy_intp len, + npy_int8 *res, + const npy_int8 *low, + const npy_int8 *hi) +{ + irk_rand_bounded_broadcast(state, len, res, + low, hi); +} + +void irk_rand_uint8_broadcast(irk_state *state, + npy_intp len, + npy_uint8 *res, + const npy_uint8 *low, + const npy_uint8 *hi) +{ + irk_rand_bounded_broadcast(state, len, + res, low, hi); +} + +void irk_rand_int16_broadcast(irk_state *state, + npy_intp len, + npy_int16 *res, + const npy_int16 *low, + const npy_int16 *hi) +{ + irk_rand_bounded_broadcast(state, len, + res, low, hi); +} + +void irk_rand_uint16_broadcast(irk_state *state, + npy_intp len, + npy_uint16 *res, + const npy_uint16 *low, + const npy_uint16 *hi) +{ + irk_rand_bounded_broadcast( + state, len, res, low, hi); +} + +void irk_rand_int32_broadcast(irk_state *state, + npy_intp len, + npy_int32 *res, + const npy_int32 *low, + const npy_int32 *hi) +{ + irk_rand_bounded_broadcast(state, len, + res, low, hi); +} + +void irk_rand_uint32_broadcast(irk_state *state, + npy_intp len, + npy_uint32 *res, + const npy_uint32 *low, + const npy_uint32 *hi) +{ + irk_rand_bounded_broadcast( + state, len, res, low, hi); +} + +void irk_rand_int64_broadcast(irk_state *state, + npy_intp len, + npy_int64 *res, + const npy_int64 *low, + const npy_int64 *hi) +{ + irk_rand_bounded_broadcast(state, len, + res, low, hi); +} + +void irk_rand_uint64_broadcast(irk_state *state, + npy_intp len, + npy_uint64 *res, + const npy_uint64 *low, + const npy_uint64 *hi) +{ + irk_rand_bounded_broadcast( + state, len, res, low, hi); +} + const MKL_INT cholesky_storage_flags[3] = {VSL_MATRIX_STORAGE_FULL, VSL_MATRIX_STORAGE_PACKED, VSL_MATRIX_STORAGE_DIAGONAL}; diff --git a/mkl_random/src/mkl_distributions.h b/mkl_random/src/mkl_distributions.h index ce4014a0..38fa288a 100644 --- a/mkl_random/src/mkl_distributions.h +++ b/mkl_random/src/mkl_distributions.h @@ -304,6 +304,52 @@ extern "C" const npy_bool lo, const npy_bool hi); + extern void irk_rand_int64_broadcast(irk_state *state, + npy_intp len, + npy_int64 *res, + const npy_int64 *low, + const npy_int64 *hi); + extern void irk_rand_uint64_broadcast(irk_state *state, + npy_intp len, + npy_uint64 *res, + const npy_uint64 *low, + const npy_uint64 *hi); + extern void irk_rand_int32_broadcast(irk_state *state, + npy_intp len, + npy_int32 *res, + const npy_int32 *low, + const npy_int32 *hi); + extern void irk_rand_uint32_broadcast(irk_state *state, + npy_intp len, + npy_uint32 *res, + const npy_uint32 *low, + const npy_uint32 *hi); + extern void irk_rand_int16_broadcast(irk_state *state, + npy_intp len, + npy_int16 *res, + const npy_int16 *low, + const npy_int16 *hi); + extern void irk_rand_uint16_broadcast(irk_state *state, + npy_intp len, + npy_uint16 *res, + const npy_uint16 *low, + const npy_uint16 *hi); + extern void irk_rand_int8_broadcast(irk_state *state, + npy_intp len, + npy_int8 *res, + const npy_int8 *low, + const npy_int8 *hi); + extern void irk_rand_uint8_broadcast(irk_state *state, + npy_intp len, + npy_uint8 *res, + const npy_uint8 *low, + const npy_uint8 *hi); + extern void irk_rand_bool_broadcast(irk_state *state, + npy_intp len, + npy_bool *res, + const npy_bool *low, + const npy_bool *hi); + extern void irk_ulong_vec(irk_state *state, npy_intp len, unsigned long *res); extern void irk_long_vec(irk_state *state, npy_intp len, long *res); From 7ebcee07883f5ae1169100212369548935707eda Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Tue, 1 Sep 2026 04:42:55 -0700 Subject: [PATCH 02/14] Add broadcast path to randint for array-like low/high --- mkl_random/mklrand.pyx | 306 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 270 insertions(+), 36 deletions(-) diff --git a/mkl_random/mklrand.pyx b/mkl_random/mklrand.pyx index edbcd400..b776a4ba 100644 --- a/mkl_random/mklrand.pyx +++ b/mkl_random/mklrand.pyx @@ -392,6 +392,70 @@ cdef extern from "mkl_distributions.h": cnp.npy_int64 high ) noexcept nogil + void irk_rand_bool_broadcast( + irk_state *state, + cnp.npy_intp len, + cnp.npy_bool *res, + const cnp.npy_bool *low, + const cnp.npy_bool *high + ) noexcept nogil + void irk_rand_uint8_broadcast( + irk_state *state, + cnp.npy_intp len, + cnp.npy_uint8 *res, + const cnp.npy_uint8 *low, + const cnp.npy_uint8 *high + ) noexcept nogil + void irk_rand_int8_broadcast( + irk_state *state, + cnp.npy_intp len, + cnp.npy_int8 *res, + const cnp.npy_int8 *low, + const cnp.npy_int8 *high + ) noexcept nogil + void irk_rand_uint16_broadcast( + irk_state *state, + cnp.npy_intp len, + cnp.npy_uint16 *res, + const cnp.npy_uint16 *low, + const cnp.npy_uint16 *high + ) noexcept nogil + void irk_rand_int16_broadcast( + irk_state *state, + cnp.npy_intp len, + cnp.npy_int16 *res, + const cnp.npy_int16 *low, + const cnp.npy_int16 *high + ) noexcept nogil + void irk_rand_uint32_broadcast( + irk_state *state, + cnp.npy_intp len, + cnp.npy_uint32 *res, + const cnp.npy_uint32 *low, + const cnp.npy_uint32 *high + ) noexcept nogil + void irk_rand_int32_broadcast( + irk_state *state, + cnp.npy_intp len, + cnp.npy_int32 *res, + const cnp.npy_int32 *low, + const cnp.npy_int32 *high + ) noexcept nogil + void irk_rand_uint64_broadcast( + irk_state *state, + cnp.npy_intp len, + cnp.npy_uint64 *res, + const cnp.npy_uint64 *low, + const cnp.npy_uint64 *high + ) noexcept nogil + void irk_rand_int64_broadcast( + irk_state *state, + cnp.npy_intp len, + cnp.npy_int64 *res, + const cnp.npy_int64 *low, + const cnp.npy_int64 *high + ) noexcept nogil + void irk_long_vec( irk_state *state, cnp.npy_intp len, long *res ) noexcept nogil @@ -1873,22 +1937,26 @@ cdef class _MKLRandomState: # a few places. It would be easy to template them. def _choose_randint_type(self, dtype): - _randint_type = { - "bool": (0, 2, self._rand_bool), - "int8": (-2**7, 2**7, self._rand_int8), - "int16": (-2**15, 2**15, self._rand_int16), - "int32": (-2**31, 2**31, self._rand_int32), - "int64": (-2**63, 2**63, self._rand_int64), - "uint8": (0, 2**8, self._rand_uint8), - "uint16": (0, 2**16, self._rand_uint16), - "uint32": (0, 2**32, self._rand_uint32), - "uint64": (0, 2**64, self._rand_uint64) + _randint_bounds = { + "bool": (0, 2), + "int8": (-2**7, 2**7), + "int16": (-2**15, 2**15), + "int32": (-2**31, 2**31), + "int64": (-2**63, 2**63), + "uint8": (0, 2**8), + "uint16": (0, 2**16), + "uint32": (0, 2**32), + "uint64": (0, 2**64), } key = np.dtype(dtype).name - if key not in _randint_type: + if key not in _randint_bounds: raise TypeError(f'Unsupported dtype "{key}" for randint') - return _randint_type[key] + + lowbnd, highbnd = _randint_bounds[key] + return (lowbnd, highbnd, + getattr(self, f"_rand_{key}"), + getattr(self, f"_rand_{key}_broadcast")) # generates typed random integer in [low, high] def _rand_bool(self, cnp.npy_bool low, cnp.npy_bool high, size): @@ -2119,6 +2187,141 @@ cdef class _MKLRandomState: irk_rand_uint64_vec(self.internal_state, cnt, out, low, high) return array + # Broadcasted variants of the typed generators for randint + def _rand_bool_broadcast(self, cnp.ndarray low, cnp.ndarray high, + cnp.ndarray out): + cdef cnp.npy_intp cnt = cnp.PyArray_SIZE(out) + cdef cnp.npy_bool *out_p = cnp.PyArray_DATA(out) + cdef cnp.npy_bool *low_p = cnp.PyArray_DATA(low) + cdef cnp.npy_bool *high_p = cnp.PyArray_DATA(high) + with nogil: + irk_rand_bool_broadcast( + self.internal_state, cnt, out_p, low_p, high_p + ) + + def _rand_int8_broadcast(self, cnp.ndarray low, cnp.ndarray high, + cnp.ndarray out): + cdef cnp.npy_intp cnt = cnp.PyArray_SIZE(out) + cdef cnp.npy_int8 *out_p = cnp.PyArray_DATA(out) + cdef cnp.npy_int8 *low_p = cnp.PyArray_DATA(low) + cdef cnp.npy_int8 *high_p = cnp.PyArray_DATA(high) + with nogil: + irk_rand_int8_broadcast( + self.internal_state, cnt, out_p, low_p, high_p + ) + + def _rand_int16_broadcast(self, cnp.ndarray low, cnp.ndarray high, + cnp.ndarray out): + cdef cnp.npy_intp cnt = cnp.PyArray_SIZE(out) + cdef cnp.npy_int16 *out_p = cnp.PyArray_DATA(out) + cdef cnp.npy_int16 *low_p = cnp.PyArray_DATA(low) + cdef cnp.npy_int16 *high_p = cnp.PyArray_DATA(high) + with nogil: + irk_rand_int16_broadcast( + self.internal_state, cnt, out_p, low_p, high_p + ) + + def _rand_int32_broadcast(self, cnp.ndarray low, cnp.ndarray high, + cnp.ndarray out): + cdef cnp.npy_intp cnt = cnp.PyArray_SIZE(out) + cdef cnp.npy_int32 *out_p = cnp.PyArray_DATA(out) + cdef cnp.npy_int32 *low_p = cnp.PyArray_DATA(low) + cdef cnp.npy_int32 *high_p = cnp.PyArray_DATA(high) + with nogil: + irk_rand_int32_broadcast( + self.internal_state, cnt, out_p, low_p, high_p + ) + + def _rand_int64_broadcast(self, cnp.ndarray low, cnp.ndarray high, + cnp.ndarray out): + cdef cnp.npy_intp cnt = cnp.PyArray_SIZE(out) + cdef cnp.npy_int64 *out_p = cnp.PyArray_DATA(out) + cdef cnp.npy_int64 *low_p = cnp.PyArray_DATA(low) + cdef cnp.npy_int64 *high_p = cnp.PyArray_DATA(high) + with nogil: + irk_rand_int64_broadcast( + self.internal_state, cnt, out_p, low_p, high_p + ) + + def _rand_uint8_broadcast(self, cnp.ndarray low, cnp.ndarray high, + cnp.ndarray out): + cdef cnp.npy_intp cnt = cnp.PyArray_SIZE(out) + cdef cnp.npy_uint8 *out_p = cnp.PyArray_DATA(out) + cdef cnp.npy_uint8 *low_p = cnp.PyArray_DATA(low) + cdef cnp.npy_uint8 *high_p = cnp.PyArray_DATA(high) + with nogil: + irk_rand_uint8_broadcast( + self.internal_state, cnt, out_p, low_p, high_p + ) + + def _rand_uint16_broadcast(self, cnp.ndarray low, cnp.ndarray high, + cnp.ndarray out): + cdef cnp.npy_intp cnt = cnp.PyArray_SIZE(out) + cdef cnp.npy_uint16 *out_p = cnp.PyArray_DATA(out) + cdef cnp.npy_uint16 *low_p = cnp.PyArray_DATA(low) + cdef cnp.npy_uint16 *high_p = cnp.PyArray_DATA(high) + with nogil: + irk_rand_uint16_broadcast( + self.internal_state, cnt, out_p, low_p, high_p + ) + + def _rand_uint32_broadcast(self, cnp.ndarray low, cnp.ndarray high, + cnp.ndarray out): + cdef cnp.npy_intp cnt = cnp.PyArray_SIZE(out) + cdef cnp.npy_uint32 *out_p = cnp.PyArray_DATA(out) + cdef cnp.npy_uint32 *low_p = cnp.PyArray_DATA(low) + cdef cnp.npy_uint32 *high_p = cnp.PyArray_DATA(high) + with nogil: + irk_rand_uint32_broadcast( + self.internal_state, cnt, out_p, low_p, high_p + ) + + def _rand_uint64_broadcast(self, cnp.ndarray low, cnp.ndarray high, + cnp.ndarray out): + cdef cnp.npy_intp cnt = cnp.PyArray_SIZE(out) + cdef cnp.npy_uint64 *out_p = cnp.PyArray_DATA(out) + cdef cnp.npy_uint64 *low_p = cnp.PyArray_DATA(low) + cdef cnp.npy_uint64 *high_p = cnp.PyArray_DATA(high) + with nogil: + irk_rand_uint64_broadcast( + self.internal_state, cnt, out_p, low_p, high_p + ) + + def _randint_broadcast(self, low_arr, high_arr, size, _dtype, + lowbnd, highbnd, broadcast_func): + # output shape + # `size` if given, else the broadcast of the bounds + if size is None: + out_shape = np.broadcast_shapes(low_arr.shape, high_arr.shape) + elif isinstance(size, (int, np.integer)): + out_shape = (int(size),) + else: + out_shape = tuple(int(s) for s in size) + + # raises ValueError if the bounds do not fit `out_shape` + low_b = np.broadcast_to(low_arr, out_shape) + high_b = np.broadcast_to(high_arr, out_shape) + + if np.prod(out_shape) == 0: + return np.empty(out_shape, dtype=_dtype) + + if int(np.min(low_b)) < lowbnd: + raise ValueError(f"low is out of bounds for {_dtype.name}") + if int(np.max(high_b)) > highbnd: + raise ValueError(f"high is out of bounds for {_dtype.name}") + if np.any(low_b >= high_b): + raise ValueError("low >= high") + + # C routine wants contiguous result-dtype arrays with `high` inclusive + low_c = np.ascontiguousarray(low_b, dtype=_dtype) + high_c = np.ascontiguousarray(high_b - 1, dtype=_dtype) + out = np.empty(out_shape, dtype=_dtype) + + with self.lock: + broadcast_func(low_c, high_c, out) + + return out + def randint(self, low, high=None, size=None, dtype=int): """ randint(low, high=None, size=None, dtype=int) @@ -2131,13 +2334,15 @@ cdef class _MKLRandomState: 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 @@ -2166,16 +2371,32 @@ cdef class _MKLRandomState: Examples -------- >>> mkl_random.randint(2, size=10) - array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0]) + array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0]) # random >>> 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], + array([[4, 0, 2, 1], # random [3, 2, 2, 0]]) + Generate a 1 x 3 array with 3 different upper bounds + + >>> mkl_random.randint(1, [3, 5, 10]) + array([2, 4, 7]) # random + + Generate a 1 by 3 array with 3 different lower bounds + + >>> mkl_random.randint([1, 5, 7], 10) + array([6, 6, 9]) # random + + Generate a 2 by 4 array using broadcasting with dtype of uint8 + + >>> mkl_random.randint([1, 3, 5, 7], [[10], [20]], dtype=numpy.uint8) + array([[ 8, 7, 7, 7], # random + [18, 17, 19, 17]], dtype=uint8) + """ if high is None: high = low @@ -2191,30 +2412,43 @@ cdef class _MKLRandomState: "ValueError", DeprecationWarning) _dtype = _dtype.newbyteorder() - if size is not None: - if (np.prod(size) == 0): - return np.empty(size, dtype=np.dtype(_dtype)) + lowbnd, highbnd, randfunc, broadcast_func = \ + self._choose_randint_type(_dtype) - lowbnd, highbnd, randfunc = self._choose_randint_type(_dtype) + low_arr = np.asarray(low) + high_arr = np.asarray(high) - if low < lowbnd: - raise ValueError( - f"low is out of bounds for {np.dtype(_dtype).name}" - ) - if high > highbnd: - raise ValueError( - f"high is out of bounds for {np.dtype(_dtype).name}" - ) - if low >= high: - raise ValueError("low >= high") + if low_arr.ndim == 0 and high_arr.ndim == 0: + # Fast path for scalar + if size is not None and np.prod(size) == 0: + return np.empty(size, dtype=_dtype) - with self.lock: - ret = randfunc(low, high - 1, size) + low = int(low) + high = int(high) - if size is None and dtype in (bool, int): - return dtype(ret) + if low < lowbnd: + raise ValueError( + f"low is out of bounds for {_dtype.name}" + ) + if high > highbnd: + raise ValueError( + f"high is out of bounds for {_dtype.name}" + ) + if low >= high: + raise ValueError("low >= high") - return ret + with self.lock: + ret = randfunc(low, high - 1, size) + + if size is None and dtype in (bool, int): + return dtype(ret) + + return ret + + # Broadcast path( at least one of `low`/`high` is array_like) + return self._randint_broadcast( + low_arr, high_arr, size, _dtype, lowbnd, highbnd, broadcast_func + ) def bytes(self, cnp.npy_intp length): """ From 79e6c2e9ad66910907d8355022642219b6aeb461 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Tue, 1 Sep 2026 04:57:18 -0700 Subject: [PATCH 03/14] Fix flake8 D102 ignore to cover tests/*.py --- .flake8 | 1 + 1 file changed, 1 insertion(+) diff --git a/.flake8 b/.flake8 index ff088a9b..a8e21cf0 100644 --- a/.flake8 +++ b/.flake8 @@ -26,6 +26,7 @@ extend-ignore = per-file-ignores = mkl_random/__init__.py: F401 mkl_random/interfaces/__init__.py: F401 + mkl_random/tests/*.py: D102 mkl_random/tests/**/*.py: D102 filename = *.py, *.pyx, *.pxi, *.pxd From 36bc4b7ac2f2fa9df2f4afaff18adb07679285ca Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Tue, 1 Sep 2026 04:57:33 -0700 Subject: [PATCH 04/14] Add TestRandint with array-like bounds tests --- mkl_random/tests/test_random.py | 243 ++++++++++++++++++++------------ 1 file changed, 150 insertions(+), 93 deletions(-) diff --git a/mkl_random/tests/test_random.py b/mkl_random/tests/test_random.py index 4e74d769..e3328fba 100644 --- a/mkl_random/tests/test_random.py +++ b/mkl_random/tests/test_random.py @@ -236,100 +236,157 @@ def randint(): return RandIntData(rfunc_method, integral_dtypes) -def test_randint_unsupported_type(randint): - pytest.raises(TypeError, randint.rfunc, 1, dtype=np.float64) - - -def test_randint_bounds_checking(randint): - for dt in randint.itype: - lbnd = 0 if dt is np.bool_ else np.iinfo(dt).min - ubnd = 2 if dt is np.bool_ else np.iinfo(dt).max + 1 - pytest.raises(ValueError, randint.rfunc, lbnd - 1, ubnd, dtype=dt) - pytest.raises(ValueError, randint.rfunc, lbnd, ubnd + 1, dtype=dt) - pytest.raises(ValueError, randint.rfunc, ubnd, lbnd, dtype=dt) - pytest.raises(ValueError, randint.rfunc, 1, 0, dtype=dt) - - -def test_randint_rng_zero_and_extremes(randint): - for dt in randint.itype: - lbnd = 0 if dt is np.bool_ else np.iinfo(dt).min - ubnd = 2 if dt is np.bool_ else np.iinfo(dt).max + 1 - tgt = ubnd - 1 - assert_equal(randint.rfunc(tgt, tgt + 1, size=1000, dtype=dt), tgt) - tgt = lbnd - assert_equal(randint.rfunc(tgt, tgt + 1, size=1000, dtype=dt), tgt) - tgt = lbnd + ((ubnd - lbnd) // 2) - assert_equal(randint.rfunc(tgt, tgt + 1, size=1000, dtype=dt), tgt) - - -def test_randint_in_bounds_fuzz(randint): - # Don't use fixed seed - rnd.seed() - for dt in randint.itype[1:]: - for ubnd in [4, 8, 16]: - vals = randint.rfunc(2, ubnd, size=2**16, dtype=dt) - assert_(vals.max() < ubnd) - assert_(vals.min() >= 2) - vals = randint.rfunc(0, 2, size=2**16, dtype="bool") - assert vals.max() < 2 - assert vals.min() >= 0 - - -def test_randint_repeatability(randint): - import hashlib - - # We use a md5 hash of generated sequences of 1000 samples - # in the range [0, 6) for all but np.bool, where the range - # is [0, 2). Hashes are for little endian numbers. - tgt = { - "bool": "4fee98a6885457da67c39331a9ec336f", - "int16": "80a5ff69c315ab6f80b03da1d570b656", - "int32": "15a3c379b6c7b0f296b162194eab68bc", - "int64": "ea9875f9334c2775b00d4976b85a1458", - "int8": "0f56333af47de94930c799806158a274", - "uint16": "80a5ff69c315ab6f80b03da1d570b656", - "uint32": "15a3c379b6c7b0f296b162194eab68bc", - "uint64": "ea9875f9334c2775b00d4976b85a1458", - "uint8": "0f56333af47de94930c799806158a274", - } - - for dt in randint.itype[1:]: +class TestRandint: + def test_unsupported_type(self, randint): + pytest.raises(TypeError, randint.rfunc, 1, dtype=np.float64) + + def test_bounds_checking(self, randint): + for dt in randint.itype: + lbnd = 0 if dt is np.bool_ else np.iinfo(dt).min + ubnd = 2 if dt is np.bool_ else np.iinfo(dt).max + 1 + pytest.raises(ValueError, randint.rfunc, lbnd - 1, ubnd, dtype=dt) + pytest.raises(ValueError, randint.rfunc, lbnd, ubnd + 1, dtype=dt) + pytest.raises(ValueError, randint.rfunc, ubnd, lbnd, dtype=dt) + pytest.raises(ValueError, randint.rfunc, 1, 0, dtype=dt) + + def test_rng_zero_and_extremes(self, randint): + for dt in randint.itype: + lbnd = 0 if dt is np.bool_ else np.iinfo(dt).min + ubnd = 2 if dt is np.bool_ else np.iinfo(dt).max + 1 + tgt = ubnd - 1 + assert_equal(randint.rfunc(tgt, tgt + 1, size=1000, dtype=dt), tgt) + tgt = lbnd + assert_equal(randint.rfunc(tgt, tgt + 1, size=1000, dtype=dt), tgt) + tgt = lbnd + ((ubnd - lbnd) // 2) + assert_equal(randint.rfunc(tgt, tgt + 1, size=1000, dtype=dt), tgt) + + def test_in_bounds_fuzz(self, randint): + # Don't use fixed seed + rnd.seed() + for dt in randint.itype[1:]: + for ubnd in [4, 8, 16]: + vals = randint.rfunc(2, ubnd, size=2**16, dtype=dt) + assert_(vals.max() < ubnd) + assert_(vals.min() >= 2) + vals = randint.rfunc(0, 2, size=2**16, dtype="bool") + assert vals.max() < 2 + assert vals.min() >= 0 + + def test_repeatability(self, randint): + import hashlib + + # We use a md5 hash of generated sequences of 1000 samples + # in the range [0, 6) for all but np.bool, where the range + # is [0, 2). Hashes are for little endian numbers. + tgt = { + "bool": "4fee98a6885457da67c39331a9ec336f", + "int16": "80a5ff69c315ab6f80b03da1d570b656", + "int32": "15a3c379b6c7b0f296b162194eab68bc", + "int64": "ea9875f9334c2775b00d4976b85a1458", + "int8": "0f56333af47de94930c799806158a274", + "uint16": "80a5ff69c315ab6f80b03da1d570b656", + "uint32": "15a3c379b6c7b0f296b162194eab68bc", + "uint64": "ea9875f9334c2775b00d4976b85a1458", + "uint8": "0f56333af47de94930c799806158a274", + } + + for dt in randint.itype[1:]: + rnd.seed(1234, brng="MT19937") + + # view as little endian for hash + if sys.byteorder == "little": + val = randint.rfunc(0, 6, size=1000, dtype=dt) + else: + val = randint.rfunc(0, 6, size=1000, dtype=dt).byteswap() + + res = hashlib.md5(val.view(np.int8)).hexdigest() + assert tgt[np.dtype(dt).name] == res + + # bools do not depend on endianness rnd.seed(1234, brng="MT19937") - - # view as little endian for hash - if sys.byteorder == "little": - val = randint.rfunc(0, 6, size=1000, dtype=dt) - else: - val = randint.rfunc(0, 6, size=1000, dtype=dt).byteswap() - - res = hashlib.md5(val.view(np.int8)).hexdigest() - assert tgt[np.dtype(dt).name] == res - - # bools do not depend on endianness - rnd.seed(1234, brng="MT19937") - val = randint.rfunc(0, 2, size=1000, dtype="bool").view(np.int8) - res = hashlib.md5(val).hexdigest() - assert tgt[np.dtype("bool").name] == res - - -def test_randint_respect_dtype_singleton(randint): - # See gh-7203 - for dt in randint.itype: - lbnd = 0 if dt is np.bool_ else np.iinfo(dt).min - ubnd = 2 if dt is np.bool_ else np.iinfo(dt).max + 1 - - sample = randint.rfunc(lbnd, ubnd, dtype=dt) - assert_equal(sample.dtype, np.dtype(dt)) - - for dt in (bool, int): - # The legacy rng uses "long" as the default integer: - lbnd = 0 if dt is bool else np.iinfo("long").min - ubnd = 2 if dt is bool else np.iinfo("long").max + 1 - - # gh-7284: Ensure that we get Python data types - sample = randint.rfunc(lbnd, ubnd, dtype=dt) - assert not hasattr(sample, "dtype") - assert type(sample) is dt + val = randint.rfunc(0, 2, size=1000, dtype="bool").view(np.int8) + res = hashlib.md5(val).hexdigest() + assert tgt[np.dtype("bool").name] == res + + def test_respect_dtype_singleton(self, randint): + # See gh-7203 + for dt in randint.itype: + lbnd = 0 if dt is np.bool_ else np.iinfo(dt).min + ubnd = 2 if dt is np.bool_ else np.iinfo(dt).max + 1 + + sample = randint.rfunc(lbnd, ubnd, dtype=dt) + assert_equal(sample.dtype, np.dtype(dt)) + + for dt in (bool, int): + # The legacy rng uses "long" as the default integer: + lbnd = 0 if dt is bool else np.iinfo("long").min + ubnd = 2 if dt is bool else np.iinfo("long").max + 1 + + # gh-7284: Ensure that we get Python data types + sample = randint.rfunc(lbnd, ubnd, dtype=dt) + assert not hasattr(sample, "dtype") + assert type(sample) is dt + + def test_array_bounds_shapes(self, randint): + for dt in randint.itype: + low = np.array([0]) + high = np.array([1]) + assert_equal(randint.rfunc(low, high, dtype=dt).shape, (1,)) + assert_equal(randint.rfunc(low[0], high, dtype=dt).shape, (1,)) + assert_equal(randint.rfunc(low, high[0], dtype=dt).shape, (1,)) + + # broadcasting of the two bounds + assert_equal(rnd.randint([0, 10, 20], [10, 20, 30]).shape, (3,)) + assert_equal(rnd.randint(0, [5, 6, 7]).shape, (3,)) + assert_equal(rnd.randint([[0], [10]], [[5], [20]]).shape, (2, 1)) + # broadcasting of the bounds with size + assert_equal(rnd.randint([0, 0], [5, 6], size=(3, 2)).shape, (3, 2)) + # empty output + assert_equal(rnd.randint([0], [10], size=0).shape, (0,)) + + def test_array_bounds_in_range(self, randint): + low = np.array([0, 10, 100, -50, 5]) + high = np.array([3, 20, 101, -40, 6]) + for dt in (np.int32, np.int64): + vals = np.stack( + [randint.rfunc(low, high, dtype=dt) for _ in range(2000)] + ) + assert np.all(vals >= low) + assert np.all(vals < high) + # a single-value range must always return that value + assert np.all(vals[:, 2] == 100) + + def test_array_bounds_full_range(self): + for dt, hi in [ + (np.uint8, 2**8), + (np.uint16, 2**16), + (np.uint32, 2**32), + (np.uint64, 2**64), + (np.int64, 2**63), + ]: + low = np.zeros( + 1000, dtype=np.int64 if dt is np.int64 else np.uint64 + ) + high = np.full(1000, hi, dtype=object) + vals = rnd.randint(low, high, dtype=dt) + assert vals.dtype == np.dtype(dt) + assert np.all(vals >= 0) + + def test_array_bounds_errors(self): + # low >= high in at least one element + assert_raises(ValueError, rnd.randint, [0, 5], [5, 5]) + # bounds out of dtype range + assert_raises(ValueError, rnd.randint, [-1], [5], None, np.uint8) + assert_raises(ValueError, rnd.randint, [0], [300], None, np.uint8) + # bounds incompatible with the requested size + assert_raises(ValueError, rnd.randint, [0, 0], [5, 6], (4,)) + + def test_array_bounds_repeatability(self): + low = [0, 10] + high = [100, 200] + a = rnd.MKLRandomState(5).randint(low, high, size=(1000, 2)) + b = rnd.MKLRandomState(5).randint(low, high, size=(1000, 2)) + assert_equal(a, b) class RandomDistData(NamedTuple): From 6fba7b8c4d85426c42a0d2c91db3d4ee187da064 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Tue, 1 Sep 2026 05:00:12 -0700 Subject: [PATCH 05/14] Unskip test_randint in test_numpy_random.py --- mkl_random/tests/third_party/test_numpy_random.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/mkl_random/tests/third_party/test_numpy_random.py b/mkl_random/tests/third_party/test_numpy_random.py index 71ec88ee..f7525e79 100644 --- a/mkl_random/tests/third_party/test_numpy_random.py +++ b/mkl_random/tests/third_party/test_numpy_random.py @@ -1082,8 +1082,6 @@ def test_two_arg_funcs(self): out = func(argOne, argTwo[0]) assert_equal(out.shape, tgtShape) - # TODO: fix randint to handle single arrays correctly, remove skip - @pytest.mark.skip("randint does not work with arrays") def test_randint(self): _, _, _, tgtShape = self._create_arrays() itype = [ From 293e5ed1843a790bfb86b7a49ba3a466cb759576 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Tue, 1 Sep 2026 05:25:40 -0700 Subject: [PATCH 06/14] Fix and cover high-1 overflow for narrow randint bounds --- mkl_random/mklrand.pyx | 12 +++++++++--- mkl_random/tests/test_random.py | 14 ++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/mkl_random/mklrand.pyx b/mkl_random/mklrand.pyx index b776a4ba..342b9b40 100644 --- a/mkl_random/mklrand.pyx +++ b/mkl_random/mklrand.pyx @@ -2305,16 +2305,22 @@ cdef class _MKLRandomState: if np.prod(out_shape) == 0: return np.empty(out_shape, dtype=_dtype) + max_high = int(np.max(high_b)) if int(np.min(low_b)) < lowbnd: raise ValueError(f"low is out of bounds for {_dtype.name}") - if int(np.max(high_b)) > highbnd: + if max_high > highbnd: raise ValueError(f"high is out of bounds for {_dtype.name}") if np.any(low_b >= high_b): raise ValueError("low >= high") - # C routine wants contiguous result-dtype arrays with `high` inclusive + # inclusive high (high - 1); widen small dtypes to int64 to avoid + # overflow, but subtract first when high exceeds int64 + if max_high <= 2**63 - 1: + high_incl = high_b.astype(np.int64) - 1 + else: + high_incl = high_b - 1 low_c = np.ascontiguousarray(low_b, dtype=_dtype) - high_c = np.ascontiguousarray(high_b - 1, dtype=_dtype) + high_c = np.ascontiguousarray(high_incl, dtype=_dtype) out = np.empty(out_shape, dtype=_dtype) with self.lock: diff --git a/mkl_random/tests/test_random.py b/mkl_random/tests/test_random.py index e3328fba..bd77d5d3 100644 --- a/mkl_random/tests/test_random.py +++ b/mkl_random/tests/test_random.py @@ -372,6 +372,20 @@ def test_array_bounds_full_range(self): assert vals.dtype == np.dtype(dt) assert np.all(vals >= 0) + def test_array_bounds_narrow_input_dtype(self, randint): + for in_dt, res_dt in [ + (np.int8, np.int64), + (np.int8, np.int32), + (np.int16, np.int64), + ]: + low = np.array([np.iinfo(res_dt).min // 2], dtype=res_dt) + high = np.array([np.iinfo(in_dt).min], dtype=in_dt) + vals = np.stack( + [randint.rfunc(low, high, dtype=res_dt) for _ in range(2000)] + ) + assert np.all(vals >= low) + assert np.all(vals < high.astype(res_dt)) + def test_array_bounds_errors(self): # low >= high in at least one element assert_raises(ValueError, rnd.randint, [0, 5], [5, 5]) From fca9f1352b5e9b2a90b5c52080b1b95fe77be729 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Tue, 1 Sep 2026 06:24:22 -0700 Subject: [PATCH 07/14] Optimize randint broadcast draws via Lemire's method --- mkl_random/src/mkl_distributions.cpp | 98 ++++++++++++++++++---------- 1 file changed, 62 insertions(+), 36 deletions(-) diff --git a/mkl_random/src/mkl_distributions.cpp b/mkl_random/src/mkl_distributions.cpp index b8311c28..c379ed74 100644 --- a/mkl_random/src/mkl_distributions.cpp +++ b/mkl_random/src/mkl_distributions.cpp @@ -2186,25 +2186,36 @@ static inline void } } -/* Smallest bit mask (2^k - 1) that is >= rng. */ -template -static inline UT irk_gen_mask(UT rng) +/* mulhi for Lemire (word * s): top 32 bits of the product, low to *lo. */ +static inline npy_uint32 irk_mulhi(npy_uint32 a, npy_uint32 b, npy_uint32 *lo) { - UT mask = rng; - unsigned int s = 0; - - for (s = 1; s < sizeof(UT) * 8; s <<= 1) - mask |= mask >> s; + npy_uint64 m = ((npy_uint64)a) * ((npy_uint64)b); + *lo = (npy_uint32)m; + return (npy_uint32)(m >> 32); +} - return mask; +/* Lemire multiplies word * s; for uint64 that spans 128 bits. Compute the + * high half via 32-bit schoolbook parts to avoid needing a 128-bit type. */ +static inline npy_uint64 irk_mulhi(npy_uint64 a, npy_uint64 b, npy_uint64 *lo) +{ + npy_uint64 a_lo = (npy_uint32)a, a_hi = a >> 32; + npy_uint64 b_lo = (npy_uint32)b, b_hi = b >> 32; + npy_uint64 t = a_lo * b_lo; + npy_uint64 w0 = (npy_uint32)t, carry = t >> 32; + t = a_hi * b_lo + carry; + npy_uint64 w1 = (npy_uint32)t, w2 = t >> 32; + t = a_lo * b_hi + w1; + *lo = (t << 32) | w0; + return a_hi * b_hi + w2 + (t >> 32); } /* - * Draw res[i] uniformly from [low[i], hi[i]] (inclusive) using per-element - * masked rejection, the same algorithm as irk_rand_uint64_vec but with - * per-element bounds. Words are generated in bulk by MKL; rejected elements - * are gathered into `idx` (allocated lazily) and retried on the next round. - * T is the result type, UT its unsigned counterpart, WT the raw-word type. + * Draw res[i] uniformly from [low[i], hi[i]] (inclusive) using Lemire's + * multiply-shift method (per-element bounds, same as NumPy). + * Words are generated in bulk by MKL; the rare rejected elements are + * gathered into `idx` (allocated lazily) and retried on the next round. + * T is the result type, UT its unsigned counterpart, + * WT the raw-word type (s wraps to 0 for a full-range draw). */ template static void irk_rand_bounded_broadcast(irk_state *state, @@ -2228,40 +2239,55 @@ static void irk_rand_bounded_broadcast(irk_state *state, irk_uniform_bits_vec(state, len, words); for (i = 0; i < len; ++i) { - UT rng = ((UT)hi[i]) - ((UT)low[i]); - UT value = ((UT)words[i]) & irk_gen_mask(rng); - - if (value <= rng) { - res[i] = (T)(((UT)low[i]) + value); - } - else { - if (idx == nullptr) { - idx = (npy_intp *)mkl_malloc(len * sizeof(npy_intp), 64); - assert(idx != nullptr); + WT w = (WT)words[i]; + WT s = (WT)(((UT)hi[i]) - ((UT)low[i])) + 1; /* 0 iff full range */ + WT result = w; + + if (s != 0) { + WT lo = 0; + result = irk_mulhi(w, s, &lo); + if (lo < s) { /* rare */ + WT t = (WT)(0 - s) % s; + if (lo < t) { + if (idx == nullptr) { + idx = + (npy_intp *)mkl_malloc(len * sizeof(npy_intp), 64); + assert(idx != nullptr); + } + idx[n_pending++] = i; + continue; + } } - idx[n_pending++] = i; } + res[i] = (T)(((UT)low[i]) + (UT)result); } while (n_pending > 0) { - npy_intp w = 0; + npy_intp wpos = 0; irk_uniform_bits_vec(state, n_pending, words); for (k = 0; k < n_pending; ++k) { npy_intp j = idx[k]; - UT rng = ((UT)hi[j]) - ((UT)low[j]); - UT value = ((UT)words[k]) & irk_gen_mask(rng); - - if (value <= rng) { - res[j] = (T)(((UT)low[j]) + value); - } - else { - /* keep this element pending; w <= k so idx[k] is read first */ - idx[w++] = j; + WT w = (WT)words[k]; + WT s = (WT)(((UT)hi[j]) - ((UT)low[j])) + 1; + WT result = w; + + if (s != 0) { + WT lo = 0; + result = irk_mulhi(w, s, &lo); + if (lo < s) { + WT t = (WT)(0 - s) % s; + if (lo < t) { + /* keep pending; wpos <= k so idx[k] read first */ + idx[wpos++] = j; + continue; + } + } } + res[j] = (T)(((UT)low[j]) + (UT)result); } - n_pending = w; + n_pending = wpos; } if (idx != nullptr) From 79ce5202eb9ab2d1742ced44236929f1cb99de07 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Tue, 1 Sep 2026 06:46:47 -0700 Subject: [PATCH 08/14] Add gh-168 to changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7319acd2..47c5b86a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 # [dev] (MM/DD/YYYY) ### Added +* Added support for `array_like` (broadcastable) `low`/`high` bounds in `randint` [gh-168](https://github.com/IntelPython/mkl_random/pull/168) ### Changed From 1b54bcbce3d9ef54aef73e2c2c327016f99de239 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Wed, 2 Sep 2026 06:32:06 -0700 Subject: [PATCH 09/14] Align randint examples with the others --- mkl_random/mklrand.pyx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mkl_random/mklrand.pyx b/mkl_random/mklrand.pyx index c5e6ea2b..be42ae02 100644 --- a/mkl_random/mklrand.pyx +++ b/mkl_random/mklrand.pyx @@ -2399,7 +2399,7 @@ cdef class _MKLRandomState: Generate a 2 by 4 array using broadcasting with dtype of uint8 - >>> mkl_random.randint([1, 3, 5, 7], [[10], [20]], dtype=numpy.uint8) + >>> mkl_random.randint([1, 3, 5, 7], [[10], [20]], dtype=np.uint8) array([[ 8, 7, 7, 7], # random [18, 17, 19, 17]], dtype=uint8) From ebae45ad84e098fd907ea04b21c8db5afede4773 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Wed, 2 Sep 2026 06:39:28 -0700 Subject: [PATCH 10/14] Fix and cover randint OOB for int8/int16 array bounds --- mkl_random/src/mkl_distributions.cpp | 7 +++++-- mkl_random/tests/test_random.py | 9 +++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/mkl_random/src/mkl_distributions.cpp b/mkl_random/src/mkl_distributions.cpp index c379ed74..bab02ba4 100644 --- a/mkl_random/src/mkl_distributions.cpp +++ b/mkl_random/src/mkl_distributions.cpp @@ -2240,7 +2240,9 @@ static void irk_rand_bounded_broadcast(irk_state *state, for (i = 0; i < len; ++i) { WT w = (WT)words[i]; - WT s = (WT)(((UT)hi[i]) - ((UT)low[i])) + 1; /* 0 iff full range */ + /* diff cast back to UT so narrow types wrap (no signed promotion) */ + UT d = (UT)(((UT)hi[i]) - ((UT)low[i])); + WT s = (WT)d + 1; /* 0 iff full range (32/64-bit only) */ WT result = w; if (s != 0) { @@ -2270,7 +2272,8 @@ static void irk_rand_bounded_broadcast(irk_state *state, for (k = 0; k < n_pending; ++k) { npy_intp j = idx[k]; WT w = (WT)words[k]; - WT s = (WT)(((UT)hi[j]) - ((UT)low[j])) + 1; + UT d = (UT)(((UT)hi[j]) - ((UT)low[j])); + WT s = (WT)d + 1; WT result = w; if (s != 0) { diff --git a/mkl_random/tests/test_random.py b/mkl_random/tests/test_random.py index ead6d23e..f3183ec0 100644 --- a/mkl_random/tests/test_random.py +++ b/mkl_random/tests/test_random.py @@ -386,6 +386,15 @@ def test_array_bounds_narrow_input_dtype(self, randint): assert np.all(vals >= low) assert np.all(vals < high.astype(res_dt)) + def test_array_bounds_narrow_dtype_negative_low(self): + for dt in [np.int8, np.int16]: + vals = np.stack( + [rnd.randint([-5], [6], dtype=dt) for _ in range(5000)] + ) + assert vals.dtype == np.dtype(dt) + assert np.all(vals >= -5) + assert np.all(vals < 6) + def test_array_bounds_errors(self): # low >= high in at least one element assert_raises(ValueError, rnd.randint, [0, 5], [5, 5]) From d514e616bfc05723a436797ac13c849a2b0b3e5c Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Wed, 2 Sep 2026 06:53:03 -0700 Subject: [PATCH 11/14] Fix randint array bounds with size=() to match NumPy --- mkl_random/mklrand.pyx | 9 ++++++--- mkl_random/tests/test_random.py | 5 +++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/mkl_random/mklrand.pyx b/mkl_random/mklrand.pyx index be42ae02..49ee8f5b 100644 --- a/mkl_random/mklrand.pyx +++ b/mkl_random/mklrand.pyx @@ -2298,9 +2298,12 @@ cdef class _MKLRandomState: else: out_shape = tuple(int(s) for s in size) - # raises ValueError if the bounds do not fit `out_shape` - low_b = np.broadcast_to(low_arr, out_shape) - high_b = np.broadcast_to(high_arr, out_shape) + # size-1 bounds may collapse into a smaller-rank `size` + bshape = np.broadcast_shapes(low_arr.shape, high_arr.shape, out_shape) + if np.prod(bshape) != np.prod(out_shape): + raise ValueError("shape mismatch: bounds cannot broadcast to size") + low_b = np.broadcast_to(low_arr, bshape) + high_b = np.broadcast_to(high_arr, bshape) if np.prod(out_shape) == 0: return np.empty(out_shape, dtype=_dtype) diff --git a/mkl_random/tests/test_random.py b/mkl_random/tests/test_random.py index f3183ec0..b7e068fc 100644 --- a/mkl_random/tests/test_random.py +++ b/mkl_random/tests/test_random.py @@ -341,6 +341,10 @@ def test_array_bounds_shapes(self, randint): assert_equal(rnd.randint([[0], [10]], [[5], [20]]).shape, (2, 1)) # broadcasting of the bounds with size assert_equal(rnd.randint([0, 0], [5, 6], size=(3, 2)).shape, (3, 2)) + # length-1 bounds with size=() + r = rnd.randint([3], [9], size=()) + assert_equal(r.shape, ()) + assert 3 <= int(r) < 9 # empty output assert_equal(rnd.randint([0], [10], size=0).shape, (0,)) @@ -403,6 +407,7 @@ def test_array_bounds_errors(self): assert_raises(ValueError, rnd.randint, [0], [300], None, np.uint8) # bounds incompatible with the requested size assert_raises(ValueError, rnd.randint, [0, 0], [5, 6], (4,)) + assert_raises(ValueError, rnd.randint, [3, 4], [9, 10], ()) def test_array_bounds_repeatability(self): low = [0, 10] From 7eac1b19d6ec2fa36b95290bca4df3c050f7d1f2 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Wed, 2 Sep 2026 06:58:03 -0700 Subject: [PATCH 12/14] Speed up randint validation by checking original bounds --- mkl_random/mklrand.pyx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mkl_random/mklrand.pyx b/mkl_random/mklrand.pyx index 49ee8f5b..48bcf613 100644 --- a/mkl_random/mklrand.pyx +++ b/mkl_random/mklrand.pyx @@ -2308,12 +2308,12 @@ cdef class _MKLRandomState: if np.prod(out_shape) == 0: return np.empty(out_shape, dtype=_dtype) - max_high = int(np.max(high_b)) - if int(np.min(low_b)) < lowbnd: + max_high = int(np.max(high_arr)) + if int(np.min(low_arr)) < lowbnd: raise ValueError(f"low is out of bounds for {_dtype.name}") if max_high > highbnd: raise ValueError(f"high is out of bounds for {_dtype.name}") - if np.any(low_b >= high_b): + if np.any(low_arr >= high_arr): raise ValueError("low >= high") # inclusive high (high - 1); widen small dtypes to int64 to avoid From 9df3c8a2c8e4204d6048a38b4beff962a6d29f53 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Wed, 2 Sep 2026 09:41:34 -0700 Subject: [PATCH 13/14] Support all BRNGs in array-bounds randint via viRngUniform fallback --- mkl_random/src/mkl_distributions.cpp | 48 ++++++++++++++++++++++++++-- mkl_random/tests/test_random.py | 25 +++++++++++++++ 2 files changed, 71 insertions(+), 2 deletions(-) diff --git a/mkl_random/src/mkl_distributions.cpp b/mkl_random/src/mkl_distributions.cpp index bab02ba4..6fec54d7 100644 --- a/mkl_random/src/mkl_distributions.cpp +++ b/mkl_random/src/mkl_distributions.cpp @@ -2160,12 +2160,32 @@ static inline void irk_uniform_bits_vec(irk_state *state, npy_intp len, npy_uint32 *buf) { int err = 0; + npy_intp i = 0; while (len > 0) { MKL_INT c = (len > MKL_INT_MAX) ? (MKL_INT)MKL_INT_MAX : (MKL_INT)len; err = viRngUniformBits32(VSL_RNG_METHOD_UNIFORMBITS32_STD, state->stream, c, (unsigned int *)buf); - assert(err == VSL_STATUS_OK); + if (err == VSL_RNG_ERROR_BRNG_NOT_SUPPORTED) { + /* viRngUniformBits32 unsupported for WH/MCG31/R250/MRG32K3A + * build words from two 16-bit viRngUniform halves */ + int *tmp = (int *)mkl_malloc(c * sizeof(int), 64); + assert(tmp != nullptr); + err = viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, state->stream, c, + tmp, 0, 65536); + assert(err == VSL_STATUS_OK); + for (i = 0; i < c; ++i) + buf[i] = (npy_uint32)tmp[i]; + err = viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, state->stream, c, + tmp, 0, 65536); + assert(err == VSL_STATUS_OK); + for (i = 0; i < c; ++i) + buf[i] |= ((npy_uint32)tmp[i]) << 16; + mkl_free(tmp); + } + else { + assert(err == VSL_STATUS_OK); + } buf += c; len -= c; } @@ -2175,12 +2195,36 @@ static inline void irk_uniform_bits_vec(irk_state *state, npy_intp len, npy_uint64 *buf) { int err = 0; + npy_intp i = 0; + int sh = 0; while (len > 0) { MKL_INT c = (len > MKL_INT_MAX) ? (MKL_INT)MKL_INT_MAX : (MKL_INT)len; err = viRngUniformBits64(VSL_RNG_METHOD_UNIFORMBITS64_STD, state->stream, c, (unsigned MKL_INT64 *)buf); - assert(err == VSL_STATUS_OK); + if (err == VSL_RNG_ERROR_BRNG_NOT_SUPPORTED) { + /* viRngUniformBits64 unsupported for WH/MCG31/R250/MRG32K3A + * build words from two 16-bit viRngUniform halves */ + int *tmp = (int *)mkl_malloc(c * sizeof(int), 64); + assert(tmp != nullptr); + for (sh = 0; sh < 64; sh += 16) { + err = viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, state->stream, c, + tmp, 0, 65536); + assert(err == VSL_STATUS_OK); + if (sh == 0) { + for (i = 0; i < c; ++i) + buf[i] = (npy_uint64)(npy_uint32)tmp[i]; + } + else { + for (i = 0; i < c; ++i) + buf[i] |= ((npy_uint64)(npy_uint32)tmp[i]) << sh; + } + } + mkl_free(tmp); + } + else { + assert(err == VSL_STATUS_OK); + } buf += c; len -= c; } diff --git a/mkl_random/tests/test_random.py b/mkl_random/tests/test_random.py index b7e068fc..2928e021 100644 --- a/mkl_random/tests/test_random.py +++ b/mkl_random/tests/test_random.py @@ -376,6 +376,31 @@ def test_array_bounds_full_range(self): assert vals.dtype == np.dtype(dt) assert np.all(vals >= 0) + def test_array_bounds_all_brngs(self): + brngs = [ + "MT19937", + "SFMT19937", + "WH", + "MT2203", + "MCG31", + "R250", + "MRG32K3A", + "MCG59", + "PHILOX4X32X10", + "ARS5", + ] + N = 50000 + R = (1 << 31) + 1 + for brng in brngs: + rs = rnd.MKLRandomState(0, brng=brng) + x = rs.randint( + np.zeros(N, np.uint32), + np.full(N, R, np.uint32), + dtype=np.uint32, + ) + assert x.min() >= 0 + assert int(x.max()) < R + def test_array_bounds_narrow_input_dtype(self, randint): for in_dt, res_dt in [ (np.int8, np.int64), From 61ae313aa8825291a71b4f17a0a921eee94bff8e Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Wed, 2 Sep 2026 10:06:37 -0700 Subject: [PATCH 14/14] Single-pass word assembly in BRNG fallback --- mkl_random/src/mkl_distributions.cpp | 66 +++++++++++++++------------- 1 file changed, 36 insertions(+), 30 deletions(-) diff --git a/mkl_random/src/mkl_distributions.cpp b/mkl_random/src/mkl_distributions.cpp index 6fec54d7..3d0d30ad 100644 --- a/mkl_random/src/mkl_distributions.cpp +++ b/mkl_random/src/mkl_distributions.cpp @@ -2152,9 +2152,9 @@ void irk_rand_int64_vec(irk_state *state, } /* - * Bulk generators of raw uniform words used by the broadcasted bounded-integer - * routines below. Overloaded on the word type so that the masked-rejection - * template can pick 32- or 64-bit words at compile time. + * Bulk source of raw uniform words for the broadcasted bounded-integer + * routines below, overloaded on the word type (32/64-bit). BRNGs that lack + * viRngUniformBits fall back to assembling words from viRngUniform. */ static inline void irk_uniform_bits_vec(irk_state *state, npy_intp len, npy_uint32 *buf) @@ -2167,20 +2167,24 @@ static inline void err = viRngUniformBits32(VSL_RNG_METHOD_UNIFORMBITS32_STD, state->stream, c, (unsigned int *)buf); if (err == VSL_RNG_ERROR_BRNG_NOT_SUPPORTED) { - /* viRngUniformBits32 unsupported for WH/MCG31/R250/MRG32K3A - * build words from two 16-bit viRngUniform halves */ - int *tmp = (int *)mkl_malloc(c * sizeof(int), 64); + /* viRngUniformBits32 unsupported for WH/MCG31/R250/MRG32K3A; + * build each word from two 16-bit viRngUniform halves */ + npy_intp total = 2 * (npy_intp)c, rem = total, off = 0; + int *tmp = (int *)mkl_malloc(total * sizeof(int), 64); assert(tmp != nullptr); - err = viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, state->stream, c, - tmp, 0, 65536); - assert(err == VSL_STATUS_OK); - for (i = 0; i < c; ++i) - buf[i] = (npy_uint32)tmp[i]; - err = viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, state->stream, c, - tmp, 0, 65536); - assert(err == VSL_STATUS_OK); + /* one call unless the count exceeds MKL_INT */ + while (rem > 0) { + MKL_INT cc = + (rem > MKL_INT_MAX) ? (MKL_INT)MKL_INT_MAX : (MKL_INT)rem; + err = viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, state->stream, + cc, tmp + off, 0, 65536); + assert(err == VSL_STATUS_OK); + off += cc; + rem -= cc; + } for (i = 0; i < c; ++i) - buf[i] |= ((npy_uint32)tmp[i]) << 16; + buf[i] = ((npy_uint32)tmp[2 * i]) | + (((npy_uint32)tmp[2 * i + 1]) << 16); mkl_free(tmp); } else { @@ -2196,30 +2200,32 @@ static inline void { int err = 0; npy_intp i = 0; - int sh = 0; while (len > 0) { MKL_INT c = (len > MKL_INT_MAX) ? (MKL_INT)MKL_INT_MAX : (MKL_INT)len; err = viRngUniformBits64(VSL_RNG_METHOD_UNIFORMBITS64_STD, state->stream, c, (unsigned MKL_INT64 *)buf); if (err == VSL_RNG_ERROR_BRNG_NOT_SUPPORTED) { - /* viRngUniformBits64 unsupported for WH/MCG31/R250/MRG32K3A - * build words from two 16-bit viRngUniform halves */ - int *tmp = (int *)mkl_malloc(c * sizeof(int), 64); + /* viRngUniformBits64 unsupported for WH/MCG31/R250/MRG32K3A; + * build each word from four 16-bit viRngUniform halves */ + npy_intp total = 4 * (npy_intp)c, rem = total, off = 0; + int *tmp = (int *)mkl_malloc(total * sizeof(int), 64); assert(tmp != nullptr); - for (sh = 0; sh < 64; sh += 16) { - err = viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, state->stream, c, - tmp, 0, 65536); + /* one call unless the count exceeds MKL_INT */ + while (rem > 0) { + MKL_INT cc = + (rem > MKL_INT_MAX) ? (MKL_INT)MKL_INT_MAX : (MKL_INT)rem; + err = viRngUniform(VSL_RNG_METHOD_UNIFORM_STD, state->stream, + cc, tmp + off, 0, 65536); assert(err == VSL_STATUS_OK); - if (sh == 0) { - for (i = 0; i < c; ++i) - buf[i] = (npy_uint64)(npy_uint32)tmp[i]; - } - else { - for (i = 0; i < c; ++i) - buf[i] |= ((npy_uint64)(npy_uint32)tmp[i]) << sh; - } + off += cc; + rem -= cc; } + for (i = 0; i < c; ++i) + buf[i] = ((npy_uint64)(npy_uint32)tmp[4 * i]) | + (((npy_uint64)(npy_uint32)tmp[4 * i + 1]) << 16) | + (((npy_uint64)(npy_uint32)tmp[4 * i + 2]) << 32) | + (((npy_uint64)(npy_uint32)tmp[4 * i + 3]) << 48); mkl_free(tmp); } else {