From fc22a8bc354f7552a41e845fefc4e497e66191a0 Mon Sep 17 00:00:00 2001 From: vchamarthi Date: Thu, 3 Sep 2026 22:35:20 -0500 Subject: [PATCH 1/2] perf: batch array-valued parameters into a single fill --- CHANGELOG.md | 2 + mkl_random/mklrand.pyx | 202 ++++++++++++++++++++++++++++---- mkl_random/tests/test_random.py | 89 ++++++++++++++ 3 files changed, 272 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0346ac3..3f9bc43 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,8 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 # [dev] (MM/DD/YYYY) ### Added +* Added tests for the array-valued parameter paths of the location and scale distributions ### Changed +* Sped up `normal`, `uniform`, `exponential`, `laplace`, `gumbel`, `logistic`, `rayleigh` and `lognormal` for array-valued parameters; streams for those paths change * 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) ### Fixed diff --git a/mkl_random/mklrand.pyx b/mkl_random/mklrand.pyx index 373d16e..cb04e41 100644 --- a/mkl_random/mklrand.pyx +++ b/mkl_random/mklrand.pyx @@ -698,6 +698,160 @@ cdef object vec_cont2_array( return arr_obj +cdef object _param_out_shape(object size, tuple param_shapes): + """Result shape for a parameterised draw, matching the per-element paths.""" + cdef object out_shape + cdef object bshape + + if size is None: + return np.broadcast_shapes(*param_shapes) + + out_shape = tuple(size) if np.iterable(size) else (size,) + try: + bshape = np.broadcast_shapes(out_shape, *param_shapes) + except ValueError: + raise ValueError("size is not compatible with inputs") + if bshape != out_shape: + raise ValueError("size is not compatible with inputs") + return out_shape + + +cdef object _fill_standard2( + irk_state *state, + irk_cont2_vec func, + object out_shape, + object lock, + double std_a, + double std_b +): + """Fill an entire request with one call, using standard parameters.""" + cdef cnp.ndarray array + cdef cnp.npy_intp n + cdef double *array_data + + array = np.empty(out_shape, np.float64) + n = cnp.PyArray_SIZE(array) + if n: + array_data = cnp.PyArray_DATA(array) + with lock, nogil: + func(state, n, array_data, std_a, std_b) + return array + + +cdef object vec_loc_scale_array( + irk_state *state, + irk_cont2_vec func, + object size, + cnp.ndarray oloc, + cnp.ndarray oscale, + object lock, + double std_a, + double std_b +): + """Draw a location and scale family with array-valued parameters. + + ``func(std_a, std_b)`` yields the standardised member, so + ``loc + scale * standardised`` is exact and needs one call per request. + """ + cdef object array + + array = _fill_standard2( + state, + func, + _param_out_shape( + size, ((oloc).shape, (oscale).shape) + ), + lock, + std_a, + std_b + ) + np.multiply(array, oscale, out=array) + np.add(array, oloc, out=array) + return array + + +cdef object vec_scale_array( + irk_state *state, + irk_cont1_vec func, + object size, + cnp.ndarray oscale, + object lock, + double std_a +): + """Draw a scale family with an array-valued scale, one call per request.""" + cdef cnp.ndarray array + cdef cnp.npy_intp n + cdef double *array_data + + array = np.empty( + _param_out_shape(size, ((oscale).shape,)), np.float64 + ) + n = cnp.PyArray_SIZE(array) + if n: + array_data = cnp.PyArray_DATA(array) + with lock, nogil: + func(state, n, array_data, std_a) + np.multiply(array, oscale, out=array) + return array + + +cdef object vec_uniform_array( + irk_state *state, + irk_cont2_vec func, + object size, + cnp.ndarray olow, + cnp.ndarray ohigh, + object lock +): + """Draw uniforms over array-valued bounds, one call per request.""" + cdef object array + + array = _fill_standard2( + state, + func, + _param_out_shape( + size, ((olow).shape, (ohigh).shape) + ), + lock, + 0.0, + 1.0 + ) + np.multiply(array, np.subtract(ohigh, olow), out=array) + np.add(array, olow, out=array) + return array + + +cdef object vec_lognormal_array( + irk_state *state, + irk_cont2_vec normal_func, + object size, + cnp.ndarray omean, + cnp.ndarray osigma, + object lock +): + """Draw lognormals with array-valued parameters, one call per request. + + Uses the normal fill: the parameters sit inside the exponential, so no + affine step applies to a standardised lognormal, but exp(mean + sigma * z) does. + """ + cdef object array + + array = _fill_standard2( + state, + normal_func, + _param_out_shape( + size, ((omean).shape, (osigma).shape) + ), + lock, + 0.0, + 1.0 + ) + np.multiply(array, osigma, out=array) + np.add(array, omean, out=array) + np.exp(array, out=array) + return array + + cdef object vec_cont3_array_sc( irk_state *state, irk_cont3_vec func, @@ -2548,7 +2702,7 @@ cdef class _MKLRandomState: if np.any(olow >= ohigh): raise ValueError("low >= high") - return vec_cont2_array( + return vec_uniform_array( self.internal_state, irk_uniform_vec, size, olow, ohigh, self.lock ) @@ -2950,28 +3104,28 @@ cdef class _MKLRandomState: method, [ICDF, BOXMULLER, BOXMULLER2], _method_alias_dict_gaussian ) if method is ICDF: - return vec_cont2_array( + return vec_loc_scale_array( self.internal_state, irk_normal_vec_ICDF, size, oloc, - oscale, self.lock + oscale, self.lock, 0.0, 1.0 ) elif method is BOXMULLER2: - return vec_cont2_array( + return vec_loc_scale_array( self.internal_state, irk_normal_vec_BM2, size, oloc, - oscale, self.lock + oscale, self.lock, 0.0, 1.0 ) else: - return vec_cont2_array( + return vec_loc_scale_array( self.internal_state, irk_normal_vec_BM1, size, oloc, - oscale, self.lock + oscale, self.lock, 0.0, 1.0 ) def beta(self, a, b, size=None): @@ -3105,8 +3259,9 @@ cdef class _MKLRandomState: if np.any(np.signbit(oscale) | (oscale == 0)): raise ValueError("scale <= 0") - return vec_cont1_array( - self.internal_state, irk_exponential_vec, size, oscale, self.lock + return vec_scale_array( + self.internal_state, irk_exponential_vec, size, oscale, self.lock, + 1.0 ) def tomaxint(self, size=None): @@ -4550,8 +4705,9 @@ cdef class _MKLRandomState: if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)): raise ValueError("scale <= 0") - return vec_cont2_array( - self.internal_state, irk_laplace_vec, size, oloc, oscale, self.lock + return vec_loc_scale_array( + self.internal_state, irk_laplace_vec, size, oloc, oscale, + self.lock, 0.0, 1.0 ) def gumbel(self, loc=0.0, scale=1.0, size=None): @@ -4690,8 +4846,9 @@ cdef class _MKLRandomState: if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)): raise ValueError("scale <= 0") - return vec_cont2_array( - self.internal_state, irk_gumbel_vec, size, oloc, oscale, self.lock + return vec_loc_scale_array( + self.internal_state, irk_gumbel_vec, size, oloc, oscale, + self.lock, 0.0, 1.0 ) def logistic(self, loc=0.0, scale=1.0, size=None): @@ -4791,13 +4948,15 @@ cdef class _MKLRandomState: if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)): raise ValueError("scale <= 0") - return vec_cont2_array( + return vec_loc_scale_array( self.internal_state, irk_logistic_vec, size, oloc, oscale, - self.lock + self.lock, + 0.0, + 1.0 ) def lognormal(self, mean=0.0, sigma=1.0, size=None, method=ICDF): @@ -4952,18 +5111,18 @@ cdef class _MKLRandomState: method, [ICDF, BOXMULLER], _method_alias_dict_gaussian_short ) if method is ICDF: - return vec_cont2_array( + return vec_lognormal_array( self.internal_state, - irk_lognormal_vec_ICDF, + irk_normal_vec_ICDF, size, omean, osigma, self.lock ) else: - return vec_cont2_array( + return vec_lognormal_array( self.internal_state, - irk_lognormal_vec_BM, + irk_normal_vec_BM1, size, omean, osigma, @@ -5045,8 +5204,9 @@ cdef class _MKLRandomState: if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)): raise ValueError("scale <= 0.0") - return vec_cont1_array( - self.internal_state, irk_rayleigh_vec, size, oscale, self.lock + return vec_scale_array( + self.internal_state, irk_rayleigh_vec, size, oscale, self.lock, + 1.0 ) def wald(self, mean, scale, size=None): diff --git a/mkl_random/tests/test_random.py b/mkl_random/tests/test_random.py index 671dcd5..2efc7aa 100644 --- a/mkl_random/tests/test_random.py +++ b/mkl_random/tests/test_random.py @@ -1163,6 +1163,95 @@ def test_uniform_array_bounds_return_ndarray(): assert arr.shape == (2,) +_LOC_SCALE_DISTS = [ + ("normal", lambda r, a, b, s: r.normal(a, b, s), 2.0, 3.0), + ("laplace", lambda r, a, b, s: r.laplace(a, b, s), 2.0, 3.0), + ("gumbel", lambda r, a, b, s: r.gumbel(a, b, s), 2.0, 3.0), + ("logistic", lambda r, a, b, s: r.logistic(a, b, s), 2.0, 3.0), + ("lognormal", lambda r, a, b, s: r.lognormal(a, b, s), 0.5, 0.75), + ("uniform", lambda r, a, b, s: r.uniform(a, b, s), 2.0, 5.0), +] + + +@pytest.mark.parametrize( + "name,draw,pa,pb", _LOC_SCALE_DISTS, ids=[d[0] for d in _LOC_SCALE_DISTS] +) +def test_two_param_array_matches_scalar(name, draw, pa, pb): + # Constant-valued arrays must agree with the scalar path. + n = 8192 + scalar = draw(rnd.MKLRandomState(1234), pa, pb, n) + arrayed = draw( + rnd.MKLRandomState(1234), np.full(n, pa), np.full(n, pb), None + ) + assert arrayed.shape == scalar.shape + np.testing.assert_allclose( + arrayed, + scalar, + rtol=1e-9, + atol=1e-9 * float(np.std(scalar)), + err_msg=f"{name}: array-parameter path disagrees with scalar path", + ) + + +@pytest.mark.parametrize( + "name,draw,pa,pb", _LOC_SCALE_DISTS, ids=[d[0] for d in _LOC_SCALE_DISTS] +) +def test_two_param_array_applies_per_element(name, draw, pa, pb): + # A scale sweep must widen the spread across the result. + n = 60000 + lo = np.full(n, pa) + hi = np.linspace(pb, pb * 4.0, n) + out = draw(rnd.MKLRandomState(99), lo, hi, None) + first, last = out[: n // 4], out[-n // 4 :] + assert np.std(last) > np.std(first), ( + f"{name}: per-element parameters do not appear to be applied" + ) + + +@pytest.mark.parametrize( + "name,draw,p", + [ + ("exponential", lambda r, a, s: r.exponential(a, s), 3.0), + ("rayleigh", lambda r, a, s: r.rayleigh(a, s), 3.0), + ], + ids=["exponential", "rayleigh"], +) +def test_one_param_array_matches_scalar(name, draw, p): + n = 8192 + scalar = draw(rnd.MKLRandomState(1234), p, n) + arrayed = draw(rnd.MKLRandomState(1234), np.full(n, p), None) + assert arrayed.shape == scalar.shape + np.testing.assert_allclose( + arrayed, + scalar, + rtol=1e-9, + atol=1e-9 * float(np.std(scalar)), + err_msg=f"{name}: array-parameter path disagrees with scalar path", + ) + + +@pytest.mark.parametrize( + "loc_shape,scale_shape,size,expected", + [ + ((7,), (), None, (7,)), + ((), (7,), None, (7,)), + ((7,), (7,), None, (7,)), + ((3, 1), (4,), None, (3, 4)), + ((4,), (4,), (3, 4), (3, 4)), + ((7,), (7,), 7, (7,)), + ], +) +def test_two_param_array_broadcast_shapes(loc_shape, scale_shape, size, expected): + loc = np.zeros(loc_shape) if loc_shape else 0.0 + scale = np.ones(scale_shape) if scale_shape else 1.0 + assert rnd.MKLRandomState(5).normal(loc, scale, size).shape == expected + + +def test_two_param_array_size_incompatible(): + with pytest.raises(ValueError): + rnd.MKLRandomState(5).normal(np.zeros(5), np.ones(5), 3) + + def test_randomdist_vonmises(randomdist): rnd.seed(randomdist.seed, brng=randomdist.brng) actual = rnd.vonmises(mu=1.23, kappa=1.54, size=(3, 2)) From c4e70de6c784d6ac84c602f9bb387003aa6c6413 Mon Sep 17 00:00:00 2001 From: vchamarthi Date: Fri, 4 Sep 2026 09:08:20 -0500 Subject: [PATCH 2/2] fix review comments --- CHANGELOG.md | 5 +++-- mkl_random/mklrand.pyx | 51 +++++++++++++++--------------------------- 2 files changed, 21 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3f9bc43..d0cbbfb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,10 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 # [dev] (MM/DD/YYYY) ### Added -* Added tests for the array-valued parameter paths of the location and scale distributions +* Added tests for the array-valued parameter paths of the location and scale distributions [gh-171](https://github.com/IntelPython/mkl_random/pull/171) ### Changed -* Sped up `normal`, `uniform`, `exponential`, `laplace`, `gumbel`, `logistic`, `rayleigh` and `lognormal` for array-valued parameters; streams for those paths change +* Sped up `normal`, `uniform`, `exponential`, `laplace`, `gumbel`, `logistic`, `rayleigh` and `lognormal` for array-valued parameters [gh-171](https://github.com/IntelPython/mkl_random/pull/171) +* The random streams for the array-valued-parameter paths of the distributions above have changed: with a fixed seed these now produce different (but equally valid) samples. Scalar-parameter paths are unaffected. [gh-171](https://github.com/IntelPython/mkl_random/pull/171) * 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) ### Fixed diff --git a/mkl_random/mklrand.pyx b/mkl_random/mklrand.pyx index cb04e41..e7a82d4 100644 --- a/mkl_random/mklrand.pyx +++ b/mkl_random/mklrand.pyx @@ -720,9 +720,7 @@ cdef object _fill_standard2( irk_state *state, irk_cont2_vec func, object out_shape, - object lock, - double std_a, - double std_b + object lock ): """Fill an entire request with one call, using standard parameters.""" cdef cnp.ndarray array @@ -734,7 +732,7 @@ cdef object _fill_standard2( if n: array_data = cnp.PyArray_DATA(array) with lock, nogil: - func(state, n, array_data, std_a, std_b) + func(state, n, array_data, 0.0, 1.0) return array @@ -744,13 +742,11 @@ cdef object vec_loc_scale_array( object size, cnp.ndarray oloc, cnp.ndarray oscale, - object lock, - double std_a, - double std_b + object lock ): """Draw a location and scale family with array-valued parameters. - ``func(std_a, std_b)`` yields the standardised member, so + ``func(0.0, 1.0)`` yields the standardised member, so ``loc + scale * standardised`` is exact and needs one call per request. """ cdef object array @@ -761,9 +757,7 @@ cdef object vec_loc_scale_array( _param_out_shape( size, ((oloc).shape, (oscale).shape) ), - lock, - std_a, - std_b + lock ) np.multiply(array, oscale, out=array) np.add(array, oloc, out=array) @@ -775,8 +769,7 @@ cdef object vec_scale_array( irk_cont1_vec func, object size, cnp.ndarray oscale, - object lock, - double std_a + object lock ): """Draw a scale family with an array-valued scale, one call per request.""" cdef cnp.ndarray array @@ -790,7 +783,7 @@ cdef object vec_scale_array( if n: array_data = cnp.PyArray_DATA(array) with lock, nogil: - func(state, n, array_data, std_a) + func(state, n, array_data, 1.0) np.multiply(array, oscale, out=array) return array @@ -812,9 +805,7 @@ cdef object vec_uniform_array( _param_out_shape( size, ((olow).shape, (ohigh).shape) ), - lock, - 0.0, - 1.0 + lock ) np.multiply(array, np.subtract(ohigh, olow), out=array) np.add(array, olow, out=array) @@ -842,9 +833,7 @@ cdef object vec_lognormal_array( _param_out_shape( size, ((omean).shape, (osigma).shape) ), - lock, - 0.0, - 1.0 + lock ) np.multiply(array, osigma, out=array) np.add(array, omean, out=array) @@ -3109,7 +3098,7 @@ cdef class _MKLRandomState: irk_normal_vec_ICDF, size, oloc, - oscale, self.lock, 0.0, 1.0 + oscale, self.lock ) elif method is BOXMULLER2: return vec_loc_scale_array( @@ -3117,7 +3106,7 @@ cdef class _MKLRandomState: irk_normal_vec_BM2, size, oloc, - oscale, self.lock, 0.0, 1.0 + oscale, self.lock ) else: return vec_loc_scale_array( @@ -3125,7 +3114,7 @@ cdef class _MKLRandomState: irk_normal_vec_BM1, size, oloc, - oscale, self.lock, 0.0, 1.0 + oscale, self.lock ) def beta(self, a, b, size=None): @@ -3260,8 +3249,7 @@ cdef class _MKLRandomState: if np.any(np.signbit(oscale) | (oscale == 0)): raise ValueError("scale <= 0") return vec_scale_array( - self.internal_state, irk_exponential_vec, size, oscale, self.lock, - 1.0 + self.internal_state, irk_exponential_vec, size, oscale, self.lock ) def tomaxint(self, size=None): @@ -4707,7 +4695,7 @@ cdef class _MKLRandomState: raise ValueError("scale <= 0") return vec_loc_scale_array( self.internal_state, irk_laplace_vec, size, oloc, oscale, - self.lock, 0.0, 1.0 + self.lock ) def gumbel(self, loc=0.0, scale=1.0, size=None): @@ -4848,7 +4836,7 @@ cdef class _MKLRandomState: raise ValueError("scale <= 0") return vec_loc_scale_array( self.internal_state, irk_gumbel_vec, size, oloc, oscale, - self.lock, 0.0, 1.0 + self.lock ) def logistic(self, loc=0.0, scale=1.0, size=None): @@ -4954,9 +4942,7 @@ cdef class _MKLRandomState: size, oloc, oscale, - self.lock, - 0.0, - 1.0 + self.lock ) def lognormal(self, mean=0.0, sigma=1.0, size=None, method=ICDF): @@ -5122,7 +5108,7 @@ cdef class _MKLRandomState: else: return vec_lognormal_array( self.internal_state, - irk_normal_vec_BM1, + irk_normal_vec_BM2, size, omean, osigma, @@ -5205,8 +5191,7 @@ cdef class _MKLRandomState: if np.any(np.signbit(oscale) | np.equal(oscale, 0.0)): raise ValueError("scale <= 0.0") return vec_scale_array( - self.internal_state, irk_rayleigh_vec, size, oscale, self.lock, - 1.0 + self.internal_state, irk_rayleigh_vec, size, oscale, self.lock ) def wald(self, mean, scale, size=None):