diff --git a/numpy/random/src/distributions/distributions.c b/numpy/random/src/distributions/distributions.c index 1241329151a9..4ce4b728584c 100644 --- a/numpy/random/src/distributions/distributions.c +++ b/numpy/random/src/distributions/distributions.c @@ -1000,6 +1000,15 @@ int64_t random_geometric(bitgen_t *bitgen_state, double p) { RAND_INT_TYPE random_zipf(bitgen_t *bitgen_state, double a) { double am1, b; + if (a >= 1025) { + /* + * If a exceeds 1025, the calculation of b will overflow and the loop + * will not terminate. It is safe to simply return 1 here, because the + * probability of generating a value greater than 1 in this case is + * less than 3e-309. + */ + return (RAND_INT_TYPE) 1; + } am1 = a - 1.0; b = pow(2.0, am1); while (1) { diff --git a/numpy/random/tests/test_generator_mt19937_regressions.py b/numpy/random/tests/test_generator_mt19937_regressions.py index d451c6acd16d..4983f45743fd 100644 --- a/numpy/random/tests/test_generator_mt19937_regressions.py +++ b/numpy/random/tests/test_generator_mt19937_regressions.py @@ -163,3 +163,10 @@ def test_geometric_tiny_prob(self): # is 0.9999999999907766, so we expect the result to be all 2**63-1. assert_array_equal(self.mt19937.geometric(p=1e-30, size=3), np.iinfo(np.int64).max) + + def test_zipf_large_parameter(self): + # Regression test for part of gh-9829: a call such as rng.zipf(10000) + # would hang. + n = 8 + sample = self.mt19937.zipf(10000, size=n) + assert_array_equal(sample, np.ones(n, dtype=np.int64))