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 diff --git a/CHANGELOG.md b/CHANGELOG.md index 0346ac3b..61e54d3d 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 * Pinned Cython in the Coverity Scan workflow so generated code stays stable between scans, and added `coverity/README.md` documenting the known Cython-boilerplate false positives and the scan review checklist [gh-164](https://github.com/IntelPython/mkl_random/pull/164) diff --git a/mkl_random/mklrand.pyx b/mkl_random/mklrand.pyx index 373d16e4..48bcf613 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,150 @@ 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) + + # 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) + + 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_arr >= high_arr): + raise ValueError("low >= high") + + # 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_incl, 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 +2343,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 +2380,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=np.uint8) + array([[ 8, 7, 7, 7], # random + [18, 17, 19, 17]], dtype=uint8) + """ if high is None: high = low @@ -2191,30 +2421,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): """ diff --git a/mkl_random/src/mkl_distributions.cpp b/mkl_random/src/mkl_distributions.cpp index 8fa67d96..3d0d30ad 100644 --- a/mkl_random/src/mkl_distributions.cpp +++ b/mkl_random/src/mkl_distributions.cpp @@ -2151,6 +2151,293 @@ void irk_rand_int64_vec(irk_state *state, res[i] = res[i] + lo; } +/* + * 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) +{ + 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); + if (err == VSL_RNG_ERROR_BRNG_NOT_SUPPORTED) { + /* 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); + /* 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[2 * i]) | + (((npy_uint32)tmp[2 * i + 1]) << 16); + mkl_free(tmp); + } + else { + 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; + npy_intp i = 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 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); + /* 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_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 { + assert(err == VSL_STATUS_OK); + } + buf += c; + len -= c; + } +} + +/* 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) +{ + npy_uint64 m = ((npy_uint64)a) * ((npy_uint64)b); + *lo = (npy_uint32)m; + return (npy_uint32)(m >> 32); +} + +/* 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 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, + 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) { + WT w = (WT)words[i]; + /* 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) { + 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; + } + } + } + res[i] = (T)(((UT)low[i]) + (UT)result); + } + + while (n_pending > 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]; + WT w = (WT)words[k]; + UT d = (UT)(((UT)hi[j]) - ((UT)low[j])); + WT s = (WT)d + 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 = wpos; + } + + 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); diff --git a/mkl_random/tests/test_random.py b/mkl_random/tests/test_random.py index 671dcd56..2928e021 100644 --- a/mkl_random/tests/test_random.py +++ b/mkl_random/tests/test_random.py @@ -236,100 +236,210 @@ 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)) + # 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,)) + + 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_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), + (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_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]) + # 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,)) + assert_raises(ValueError, rnd.randint, [3, 4], [9, 10], ()) + + 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): 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 = [