diff --git a/pytensor_distributions/asymmetriclaplace.py b/pytensor_distributions/asymmetriclaplace.py index 993c5ec..2e053ca 100644 --- a/pytensor_distributions/asymmetriclaplace.py +++ b/pytensor_distributions/asymmetriclaplace.py @@ -111,7 +111,9 @@ def isf(x, mu, b, kappa): def rvs(mu, b, kappa, size=None, random_state=None): - random_samples = pt.random.uniform(-kappa, 1 / kappa, size=size, rng=random_state) + random_samples = pt.random.uniform( + -kappa, 1 / kappa, size=size, rng=random_state, return_next_rng=True + )[1] sgn = pt.sign(random_samples) return mu - (1 / (1 / b * sgn * pt.power(kappa, sgn))) * pt.log( 1 - random_samples * sgn * pt.power(kappa, sgn) diff --git a/pytensor_distributions/bernoulli.py b/pytensor_distributions/bernoulli.py index b79e5eb..7831dc9 100644 --- a/pytensor_distributions/bernoulli.py +++ b/pytensor_distributions/bernoulli.py @@ -1,5 +1,5 @@ import pytensor.tensor as pt -from pytensor.tensor.xlogx import xlogx +from pytensor.tensor.special import xlogy from pytensor_distributions.helper import cdf_bounds, ppf_bounds_disc @@ -51,11 +51,11 @@ def lmoment4(p): def entropy(p): q = 1 - p - return -xlogx(p) - xlogx(q) + return -xlogy(p, p) - xlogy(q, q) def rvs(p, size=None, random_state=None): - return pt.random.binomial(1, p, size=size, rng=random_state) + return pt.random.binomial(1, p, size=size, rng=random_state, return_next_rng=True)[1] def cdf(x, p): diff --git a/pytensor_distributions/beta.py b/pytensor_distributions/beta.py index 223931a..311151b 100644 --- a/pytensor_distributions/beta.py +++ b/pytensor_distributions/beta.py @@ -100,7 +100,7 @@ def sf(x, alpha, beta): def rvs(alpha, beta, size=None, random_state=None): - return pt.random.beta(alpha, beta, rng=random_state, size=size) + return pt.random.beta(alpha, beta, rng=random_state, size=size, return_next_rng=True)[1] def logcdf(x, alpha, beta): diff --git a/pytensor_distributions/betabinomial.py b/pytensor_distributions/betabinomial.py index d67d9eb..637721c 100644 --- a/pytensor_distributions/betabinomial.py +++ b/pytensor_distributions/betabinomial.py @@ -126,7 +126,7 @@ def isf(x, n, alpha, beta): def rvs(n, alpha, beta, size=None, random_state=None): - return pt.random.betabinom(n, alpha, beta, size=size, rng=random_state) + return pt.random.betabinom(n, alpha, beta, size=size, rng=random_state, return_next_rng=True)[1] def logcdf(x, n, alpha, beta): diff --git a/pytensor_distributions/betascaled.py b/pytensor_distributions/betascaled.py index 25193bf..24262d6 100644 --- a/pytensor_distributions/betascaled.py +++ b/pytensor_distributions/betascaled.py @@ -106,7 +106,7 @@ def sf(x, alpha, beta, lower, upper): def rvs(alpha, beta, lower, upper, size=None, random_state=None): - beta_samples = pt.random.beta(alpha, beta, rng=random_state, size=size) + beta_samples = pt.random.beta(alpha, beta, rng=random_state, size=size, return_next_rng=True)[1] return beta_samples * (upper - lower) + lower diff --git a/pytensor_distributions/binomial.py b/pytensor_distributions/binomial.py index 80d62cd..84e9625 100644 --- a/pytensor_distributions/binomial.py +++ b/pytensor_distributions/binomial.py @@ -77,7 +77,7 @@ def isf(q, n, p): def rvs(n, p, size=None, random_state=None): - return pt.random.binomial(n, p, size=size, rng=random_state) + return pt.random.binomial(n, p, size=size, rng=random_state, return_next_rng=True)[1] def logpdf(x, n, p): diff --git a/pytensor_distributions/categorical.py b/pytensor_distributions/categorical.py index 0cf0a7d..607f594 100644 --- a/pytensor_distributions/categorical.py +++ b/pytensor_distributions/categorical.py @@ -1,5 +1,5 @@ import pytensor.tensor as pt -from pytensor.tensor.xlogx import xlogx +from pytensor.tensor.special import xlogy from pytensor_distributions.helper import cdf_bounds, ppf_bounds_disc @@ -54,7 +54,7 @@ def kurtosis(p): def entropy(p): p = _normalize_p(p) - return -pt.sum(xlogx(p), axis=-1) + return -pt.sum(xlogy(p, p), axis=-1) def pdf(x, p): @@ -114,4 +114,4 @@ def isf(q, p): def rvs(p, size=None, random_state=None): p = _normalize_p(p) - return pt.random.categorical(p, size=size, rng=random_state) + return pt.random.categorical(p, size=size, rng=random_state, return_next_rng=True)[1] diff --git a/pytensor_distributions/cauchy.py b/pytensor_distributions/cauchy.py index d45e819..aece68b 100644 --- a/pytensor_distributions/cauchy.py +++ b/pytensor_distributions/cauchy.py @@ -82,7 +82,7 @@ def sf(x, alpha, beta): def rvs(alpha, beta, size=None, random_state=None): - return pt.random.cauchy(alpha, beta, rng=random_state, size=size) + return pt.random.cauchy(alpha, beta, rng=random_state, size=size, return_next_rng=True)[1] def logcdf(x, alpha, beta): diff --git a/pytensor_distributions/chisquared.py b/pytensor_distributions/chisquared.py index ec98cf4..1fffd97 100644 --- a/pytensor_distributions/chisquared.py +++ b/pytensor_distributions/chisquared.py @@ -68,7 +68,7 @@ def pdf(x, nu): def rvs(nu, size=None, random_state=None): - return pt.random.chisquare(nu, rng=random_state, size=size) + return pt.random.chisquare(nu, rng=random_state, size=size, return_next_rng=True)[1] def logcdf(x, nu): diff --git a/pytensor_distributions/dirichlet.py b/pytensor_distributions/dirichlet.py index a2637fe..0b6a2a5 100644 --- a/pytensor_distributions/dirichlet.py +++ b/pytensor_distributions/dirichlet.py @@ -73,4 +73,4 @@ def logpdf(x, alpha): def rvs(alpha, size=None, random_state=None): - return pt.random.dirichlet(alpha, size=size, rng=random_state) + return pt.random.dirichlet(alpha, size=size, rng=random_state, return_next_rng=True)[1] diff --git a/pytensor_distributions/dirichletmultinomial.py b/pytensor_distributions/dirichletmultinomial.py index 2cff9d1..c324d1b 100644 --- a/pytensor_distributions/dirichletmultinomial.py +++ b/pytensor_distributions/dirichletmultinomial.py @@ -75,5 +75,5 @@ def logpdf(x, n, a): def rvs(n, a, size=None, random_state=None): - p = pt.random.dirichlet(a, size=size, rng=random_state) - return pt.random.multinomial(n, p, rng=random_state) + next_rng, p = pt.random.dirichlet(a, size=size, rng=random_state, return_next_rng=True) + return pt.random.multinomial(n, p, rng=next_rng, return_next_rng=True)[1] diff --git a/pytensor_distributions/discreteuniform.py b/pytensor_distributions/discreteuniform.py index 9a2dfe7..bb0114f 100644 --- a/pytensor_distributions/discreteuniform.py +++ b/pytensor_distributions/discreteuniform.py @@ -81,7 +81,9 @@ def isf(q, lower, upper): def rvs(lower, upper, size=None, random_state=None): - return pt.random.integers(lower, upper + 1, size=size, rng=random_state) + return pt.random.integers(lower, upper + 1, size=size, rng=random_state, return_next_rng=True)[ + 1 + ] def logpdf(x, lower, upper): diff --git a/pytensor_distributions/discreteweibull.py b/pytensor_distributions/discreteweibull.py index 99eb919..f4a849c 100644 --- a/pytensor_distributions/discreteweibull.py +++ b/pytensor_distributions/discreteweibull.py @@ -76,7 +76,9 @@ def sf(x, q, beta): def rvs(q, beta, size=None, random_state=None): - return ppf(pt.random.uniform(0, 1, rng=random_state, size=size), q, beta) + return ppf( + pt.random.uniform(0, 1, rng=random_state, size=size, return_next_rng=True)[1], q, beta + ) def logcdf(x, q, beta): diff --git a/pytensor_distributions/exgaussian.py b/pytensor_distributions/exgaussian.py index 64f4a7d..bac41c5 100644 --- a/pytensor_distributions/exgaussian.py +++ b/pytensor_distributions/exgaussian.py @@ -88,8 +88,8 @@ def sf(x, mu, sigma, nu): def rvs(mu, sigma, nu, size=None, random_state=None): - next_rng, exp_rvs = pt.random.exponential(nu, rng=random_state, size=size).owner.outputs - normal_rvs = pt.random.normal(mu, sigma, rng=next_rng, size=size) + next_rng, exp_rvs = pt.random.exponential(nu, rng=random_state, size=size, return_next_rng=True) + normal_rvs = pt.random.normal(mu, sigma, rng=next_rng, size=size, return_next_rng=True)[1] return normal_rvs + exp_rvs diff --git a/pytensor_distributions/exponential.py b/pytensor_distributions/exponential.py index 4a8525b..5fbd9ea 100644 --- a/pytensor_distributions/exponential.py +++ b/pytensor_distributions/exponential.py @@ -80,7 +80,7 @@ def sf(x, lam): def rvs(lam, size=None, random_state=None): - return pt.random.exponential(1.0 / lam, rng=random_state, size=size) + return pt.random.exponential(1.0 / lam, rng=random_state, size=size, return_next_rng=True)[1] def logcdf(x, lam): diff --git a/pytensor_distributions/gamma.py b/pytensor_distributions/gamma.py index 84135d1..71a1b45 100644 --- a/pytensor_distributions/gamma.py +++ b/pytensor_distributions/gamma.py @@ -76,7 +76,9 @@ def isf(x, alpha, beta): def rvs(alpha, beta, size=None, random_state=None): - return pt.random.gamma(shape=alpha, scale=1 / beta, rng=random_state, size=size) + return pt.random.gamma( + shape=alpha, scale=1 / beta, rng=random_state, size=size, return_next_rng=True + )[1] def logcdf(x, alpha, beta): diff --git a/pytensor_distributions/geometric.py b/pytensor_distributions/geometric.py index 5770c8e..d614cfa 100644 --- a/pytensor_distributions/geometric.py +++ b/pytensor_distributions/geometric.py @@ -78,7 +78,7 @@ def isf(q, p): def rvs(p, size=None, random_state=None): - return pt.random.geometric(p, size=size, rng=random_state) + return pt.random.geometric(p, size=size, rng=random_state, return_next_rng=True)[1] def logpdf(x, p): diff --git a/pytensor_distributions/gumbel.py b/pytensor_distributions/gumbel.py index e159995..853a1cc 100644 --- a/pytensor_distributions/gumbel.py +++ b/pytensor_distributions/gumbel.py @@ -86,7 +86,9 @@ def sf(x, mu, beta): def rvs(mu, beta, size=None, random_state=None): - return ppf(pt.random.uniform(0, 1, rng=random_state, size=size), mu, beta) + return ppf( + pt.random.uniform(0, 1, rng=random_state, size=size, return_next_rng=True)[1], mu, beta + ) def logcdf(x, mu, beta): diff --git a/pytensor_distributions/halfcauchy.py b/pytensor_distributions/halfcauchy.py index d1f119a..e3396df 100644 --- a/pytensor_distributions/halfcauchy.py +++ b/pytensor_distributions/halfcauchy.py @@ -81,7 +81,7 @@ def sf(x, beta): def rvs(beta, size=None, random_state=None): - uniform_samples = pt.random.uniform(0, 1, rng=random_state, size=size) + uniform_samples = pt.random.uniform(0, 1, rng=random_state, size=size, return_next_rng=True)[1] return beta * pt.tan(pt.pi / 2 * uniform_samples) diff --git a/pytensor_distributions/halfnormal.py b/pytensor_distributions/halfnormal.py index 3e6ae69..b08c461 100644 --- a/pytensor_distributions/halfnormal.py +++ b/pytensor_distributions/halfnormal.py @@ -79,7 +79,7 @@ def sf(x, sigma): def rvs(sigma, size=None, random_state=None): - return pt.abs(pt.random.normal(0, sigma, rng=random_state, size=size)) + return pt.abs(pt.random.normal(0, sigma, rng=random_state, size=size, return_next_rng=True)[1]) def logpdf(x, sigma): diff --git a/pytensor_distributions/halfstudentt.py b/pytensor_distributions/halfstudentt.py index 14f6e2a..507b029 100644 --- a/pytensor_distributions/halfstudentt.py +++ b/pytensor_distributions/halfstudentt.py @@ -123,7 +123,7 @@ def sf(x, nu, sigma): def rvs(nu, sigma, size=None, random_state=None): - t_samples = pt.random.t(nu, rng=random_state, size=size) + t_samples = pt.random.t(nu, rng=random_state, size=size, return_next_rng=True)[1] return pt.abs(t_samples * sigma) diff --git a/pytensor_distributions/hypergeometric.py b/pytensor_distributions/hypergeometric.py index bbdb468..6dcc96e 100644 --- a/pytensor_distributions/hypergeometric.py +++ b/pytensor_distributions/hypergeometric.py @@ -130,4 +130,6 @@ def isf(q, N, k, n): def rvs(N, k, n, size=None, random_state=None): - return pt.random.hypergeometric(k, N - k, n, size=size, rng=random_state) + return pt.random.hypergeometric(k, N - k, n, size=size, rng=random_state, return_next_rng=True)[ + 1 + ] diff --git a/pytensor_distributions/inversegamma.py b/pytensor_distributions/inversegamma.py index 2fba0e0..6cc4623 100644 --- a/pytensor_distributions/inversegamma.py +++ b/pytensor_distributions/inversegamma.py @@ -83,7 +83,7 @@ def sf(x, alpha, beta): def rvs(alpha, beta, size=None, random_state=None): - return 1.0 / pt.random.gamma(alpha, beta, rng=random_state, size=size) + return 1.0 / pt.random.gamma(alpha, beta, rng=random_state, size=size, return_next_rng=True)[1] def logcdf(x, alpha, beta): diff --git a/pytensor_distributions/kumaraswamy.py b/pytensor_distributions/kumaraswamy.py index bfaa06c..5bbb8b7 100644 --- a/pytensor_distributions/kumaraswamy.py +++ b/pytensor_distributions/kumaraswamy.py @@ -103,7 +103,7 @@ def isf(x, a, b): def rvs(a, b, size=None, random_state=None): - u = pt.random.uniform(0, 1, size=size, rng=random_state) + u = pt.random.uniform(0, 1, size=size, rng=random_state, return_next_rng=True)[1] return ppf(u, a, b) diff --git a/pytensor_distributions/laplace.py b/pytensor_distributions/laplace.py index 586255d..4f5b13a 100644 --- a/pytensor_distributions/laplace.py +++ b/pytensor_distributions/laplace.py @@ -85,7 +85,7 @@ def isf(q, mu, b): def rvs(mu, b, size=None, random_state=None): - return pt.random.laplace(mu, b, rng=random_state, size=size) + return pt.random.laplace(mu, b, rng=random_state, size=size, return_next_rng=True)[1] def logpdf(x, mu, b): diff --git a/pytensor_distributions/logistic.py b/pytensor_distributions/logistic.py index 01baa38..2c0e6e7 100644 --- a/pytensor_distributions/logistic.py +++ b/pytensor_distributions/logistic.py @@ -83,7 +83,7 @@ def sf(x, mu, s): def rvs(mu, s, size=None, random_state=None): - return pt.random.logistic(mu, s, rng=random_state, size=size) + return pt.random.logistic(mu, s, rng=random_state, size=size, return_next_rng=True)[1] def logcdf(x, mu, s): diff --git a/pytensor_distributions/logitnormal.py b/pytensor_distributions/logitnormal.py index f863c39..ba7ca94 100644 --- a/pytensor_distributions/logitnormal.py +++ b/pytensor_distributions/logitnormal.py @@ -204,4 +204,6 @@ def isf(q, mu, sigma): def rvs(mu, sigma, size=None, random_state=None): - return pt.sigmoid(pt.random.normal(mu, sigma, rng=random_state, size=size)) + return pt.sigmoid( + pt.random.normal(mu, sigma, rng=random_state, size=size, return_next_rng=True)[1] + ) diff --git a/pytensor_distributions/loglogistic.py b/pytensor_distributions/loglogistic.py index dd48614..588ffc0 100644 --- a/pytensor_distributions/loglogistic.py +++ b/pytensor_distributions/loglogistic.py @@ -133,5 +133,5 @@ def isf(q, alpha, beta): def rvs(alpha, beta, size=None, random_state=None): - u = pt.random.uniform(size=size, rng=random_state) + u = pt.random.uniform(size=size, rng=random_state, return_next_rng=True)[1] return alpha * (u / (1 - u)) ** (1 / beta) diff --git a/pytensor_distributions/lognormal.py b/pytensor_distributions/lognormal.py index 4e94db0..6a3d014 100644 --- a/pytensor_distributions/lognormal.py +++ b/pytensor_distributions/lognormal.py @@ -81,7 +81,7 @@ def sf(x, mu, sigma): def rvs(mu, sigma, size=None, random_state=None): - return pt.exp(pt.random.normal(mu, sigma, rng=random_state, size=size)) + return pt.exp(pt.random.normal(mu, sigma, rng=random_state, size=size, return_next_rng=True)[1]) def logcdf(x, mu, sigma): diff --git a/pytensor_distributions/matrixnormal.py b/pytensor_distributions/matrixnormal.py index dacc394..36a4ee3 100644 --- a/pytensor_distributions/matrixnormal.py +++ b/pytensor_distributions/matrixnormal.py @@ -120,5 +120,5 @@ def rvs(mu, rowcov, colcov, size=None, random_state=None): else: full_shape = base_shape - Z = pt.random.normal(0, 1, size=full_shape, rng=random_state) + Z = pt.random.normal(0, 1, size=full_shape, rng=random_state, return_next_rng=True)[1] return target + L_row @ Z @ pt.swapaxes(L_col, -1, -2) diff --git a/pytensor_distributions/moyal.py b/pytensor_distributions/moyal.py index d927ab7..e9372dc 100644 --- a/pytensor_distributions/moyal.py +++ b/pytensor_distributions/moyal.py @@ -80,7 +80,7 @@ def isf(x, mu, sigma): def rvs(mu, sigma, size=None, random_state=None): - u = pt.random.uniform(size=size, rng=random_state) + u = pt.random.uniform(size=size, rng=random_state, return_next_rng=True)[1] return ppf(u, mu, sigma) diff --git a/pytensor_distributions/multinomial.py b/pytensor_distributions/multinomial.py index 77952a6..7785895 100644 --- a/pytensor_distributions/multinomial.py +++ b/pytensor_distributions/multinomial.py @@ -55,4 +55,4 @@ def logpdf(x, n, p): def rvs(n, p, size=None, random_state=None): - return pt.random.multinomial(n, p, size=size, rng=random_state) + return pt.random.multinomial(n, p, size=size, rng=random_state, return_next_rng=True)[1] diff --git a/pytensor_distributions/mvnormal.py b/pytensor_distributions/mvnormal.py index aa0b343..d71347a 100644 --- a/pytensor_distributions/mvnormal.py +++ b/pytensor_distributions/mvnormal.py @@ -85,4 +85,6 @@ def logpdf(x, mu, cov): def rvs(mu, cov, size=None, random_state=None): mu = pt.broadcast_to(mu, cov.shape[:-1]) - return pt.random.multivariate_normal(mu, cov, size=size, rng=random_state) + return pt.random.multivariate_normal( + mu, cov, size=size, rng=random_state, return_next_rng=True + )[1] diff --git a/pytensor_distributions/mvstudentt.py b/pytensor_distributions/mvstudentt.py index fe1eca3..01a9c84 100644 --- a/pytensor_distributions/mvstudentt.py +++ b/pytensor_distributions/mvstudentt.py @@ -70,8 +70,12 @@ def logpdf(x, nu, mu, cov): def rvs(nu, mu, cov, size=None, random_state=None): mu = pt.broadcast_to(mu, cov.shape[:-1]) - z = pt.random.multivariate_normal(pt.zeros_like(mu), cov, size=size, rng=random_state) - chi2 = pt.random.chisquare(nu, size=size if size is not None else 1, rng=random_state) + z = pt.random.multivariate_normal( + pt.zeros_like(mu), cov, size=size, rng=random_state, return_next_rng=True + )[1] + chi2 = pt.random.chisquare( + nu, size=size if size is not None else 1, rng=random_state, return_next_rng=True + )[1] if size is None: chi2 = chi2[None] return mu + z / pt.sqrt(chi2[..., None] / nu) diff --git a/pytensor_distributions/negativebinomial.py b/pytensor_distributions/negativebinomial.py index 642a05c..d89174a 100644 --- a/pytensor_distributions/negativebinomial.py +++ b/pytensor_distributions/negativebinomial.py @@ -81,7 +81,7 @@ def isf(q, n, p): def rvs(n, p, size=None, random_state=None): - return pt.random.negative_binomial(n, p, rng=random_state, size=size) + return pt.random.negative_binomial(n, p, rng=random_state, size=size, return_next_rng=True)[1] def logpdf(x, n, p): diff --git a/pytensor_distributions/normal.py b/pytensor_distributions/normal.py index 83d7822..9d516f6 100644 --- a/pytensor_distributions/normal.py +++ b/pytensor_distributions/normal.py @@ -85,7 +85,7 @@ def sf(x, mu, sigma): def rvs(mu, sigma, size=None, random_state=None): - return pt.random.normal(mu, sigma, rng=random_state, size=size) + return pt.random.normal(mu, sigma, rng=random_state, size=size, return_next_rng=True)[1] def logcdf(x, mu, sigma): diff --git a/pytensor_distributions/pareto.py b/pytensor_distributions/pareto.py index c4b50f7..74ef846 100644 --- a/pytensor_distributions/pareto.py +++ b/pytensor_distributions/pareto.py @@ -86,7 +86,7 @@ def isf(x, alpha, m): def rvs(alpha, m, size=None, random_state=None): - u = pt.random.uniform(size=size, rng=random_state) + u = pt.random.uniform(size=size, rng=random_state, return_next_rng=True)[1] return m / (1 - u) ** (1 / alpha) diff --git a/pytensor_distributions/poisson.py b/pytensor_distributions/poisson.py index 9733362..5240bb4 100644 --- a/pytensor_distributions/poisson.py +++ b/pytensor_distributions/poisson.py @@ -81,7 +81,7 @@ def isf(q, mu): def rvs(mu, size=None, random_state=None): - return pt.random.poisson(mu, rng=random_state, size=size) + return pt.random.poisson(mu, rng=random_state, size=size, return_next_rng=True)[1] def logpdf(x, mu): diff --git a/pytensor_distributions/polyagamma.py b/pytensor_distributions/polyagamma.py index 9a4828f..1ab0570 100644 --- a/pytensor_distributions/polyagamma.py +++ b/pytensor_distributions/polyagamma.py @@ -197,7 +197,9 @@ def rvs(h, z, size=None, random_state=None): else: gamma_size = (K, *size) - gamma_draws = pt.random.gamma(h, scale=1.0, size=gamma_size, rng=random_state) + gamma_draws = pt.random.gamma( + h, scale=1.0, size=gamma_size, rng=random_state, return_next_rng=True + )[1] z2_term = z**2 / (4 * pt.pi**2) k_bc = k.reshape((-1,) + (1,) * (gamma_draws.ndim - 1)) diff --git a/pytensor_distributions/rice.py b/pytensor_distributions/rice.py index b6682a6..bae5785 100644 --- a/pytensor_distributions/rice.py +++ b/pytensor_distributions/rice.py @@ -127,8 +127,8 @@ def isf(q, nu, sigma): def rvs(nu, sigma, size=None, random_state=None): - next_rng, x = pt.random.normal(nu, sigma, size=size, rng=random_state).owner.outputs - y = pt.random.normal(0, sigma, size=size, rng=next_rng) + next_rng, x = pt.random.normal(nu, sigma, size=size, rng=random_state, return_next_rng=True) + next_rng, y = pt.random.normal(0, sigma, size=size, rng=next_rng, return_next_rng=True) return pt.sqrt(x**2 + y**2) diff --git a/pytensor_distributions/scaledinversechisquared.py b/pytensor_distributions/scaledinversechisquared.py index 3009ec6..dc08c83 100644 --- a/pytensor_distributions/scaledinversechisquared.py +++ b/pytensor_distributions/scaledinversechisquared.py @@ -81,7 +81,9 @@ def sf(x, nu, tau2): def rvs(nu, tau2, size=None, random_state=None): - return (nu * tau2) / pt.random.chisquare(nu, rng=random_state, size=size) + return (nu * tau2) / pt.random.chisquare(nu, rng=random_state, size=size, return_next_rng=True)[ + 1 + ] def logcdf(x, nu, tau2): diff --git a/pytensor_distributions/skew_studentt.py b/pytensor_distributions/skew_studentt.py index 1d3a1a8..d707e3c 100644 --- a/pytensor_distributions/skew_studentt.py +++ b/pytensor_distributions/skew_studentt.py @@ -177,5 +177,5 @@ def isf(q, a, b, mu, sigma): def rvs(a, b, mu, sigma, size=None, random_state=None): - u = pt.random.uniform(0, 1, size=size, rng=random_state) + u = pt.random.uniform(0, 1, size=size, rng=random_state, return_next_rng=True)[1] return ppf(u, a, b, mu, sigma) diff --git a/pytensor_distributions/skewnormal.py b/pytensor_distributions/skewnormal.py index 2d60897..4fcbd2e 100644 --- a/pytensor_distributions/skewnormal.py +++ b/pytensor_distributions/skewnormal.py @@ -16,7 +16,7 @@ def mode(mu, sigma, alpha): term1 = pt.sqrt(2 / pt.pi) * delta term2 = (1 - pt.pi / 4) * (term1**3) / (1 - (2 / pt.pi) * delta**2) - term3 = 0.5 * pt.sgn(alpha) * pt.exp(-2 * pt.pi / pt.abs(alpha)) + term3 = 0.5 * pt.sign(alpha) * pt.exp(-2 * pt.pi / pt.abs(alpha)) mo_alpha = term1 - term2 - term3 return mu + sigma * mo_alpha @@ -104,8 +104,8 @@ def sf(x, mu, sigma, alpha): def rvs(mu, sigma, alpha, size=None, random_state=None): mu_b, sigma_b, alpha_b = pt.broadcast_arrays(mu, sigma, alpha) - next_rng, u_0 = pt.random.normal(0, 1, rng=random_state, size=size).owner.outputs - v = pt.random.normal(0, 1, rng=next_rng, size=size) + next_rng, u_0 = pt.random.normal(0, 1, rng=random_state, size=size, return_next_rng=True) + next_rng, v = pt.random.normal(0, 1, rng=next_rng, size=size, return_next_rng=True) d = alpha_b / pt.sqrt(1 + alpha_b**2) u_1 = d * u_0 + v * pt.sqrt(1 - d**2) return pt.sign(u_0) * u_1 * sigma_b + mu_b diff --git a/pytensor_distributions/studentt.py b/pytensor_distributions/studentt.py index 703891b..21cff8e 100644 --- a/pytensor_distributions/studentt.py +++ b/pytensor_distributions/studentt.py @@ -112,7 +112,7 @@ def sf(x, nu, mu, sigma): def rvs(nu, mu, sigma, size=None, random_state=None): - return pt.random.t(nu, mu, sigma, rng=random_state, size=size) + return pt.random.t(nu, mu, sigma, rng=random_state, size=size, return_next_rng=True)[1] def logcdf(x, nu, mu, sigma): diff --git a/pytensor_distributions/triangular.py b/pytensor_distributions/triangular.py index 3f3f575..2791350 100644 --- a/pytensor_distributions/triangular.py +++ b/pytensor_distributions/triangular.py @@ -102,7 +102,7 @@ def isf(q, lower, c, upper): def rvs(lower, c, upper, size=None, random_state=None): - u = pt.random.uniform(0.0, 1.0, size=size, rng=random_state) + u = pt.random.uniform(0.0, 1.0, size=size, rng=random_state, return_next_rng=True)[1] return ppf(u, lower, c, upper) diff --git a/pytensor_distributions/truncatednormal.py b/pytensor_distributions/truncatednormal.py index cf239d1..8cecb6d 100644 --- a/pytensor_distributions/truncatednormal.py +++ b/pytensor_distributions/truncatednormal.py @@ -274,5 +274,5 @@ def isf(q, mu, sigma, lower, upper): def rvs(mu, sigma, lower, upper, size=None, random_state=None): - u = pt.random.uniform(0, 1, size=size, rng=random_state) + u = pt.random.uniform(0, 1, size=size, rng=random_state, return_next_rng=True)[1] return ppf(u, mu, sigma, lower, upper) diff --git a/pytensor_distributions/uniform.py b/pytensor_distributions/uniform.py index 06c8e66..7eb090d 100644 --- a/pytensor_distributions/uniform.py +++ b/pytensor_distributions/uniform.py @@ -82,7 +82,7 @@ def isf(x, lower, upper): def rvs(lower, upper, size=None, random_state=None): - return pt.random.uniform(lower, upper, size=size, rng=random_state) + return pt.random.uniform(lower, upper, size=size, rng=random_state, return_next_rng=True)[1] def logcdf(x, lower, upper): diff --git a/pytensor_distributions/vonmises.py b/pytensor_distributions/vonmises.py index 0f2ea79..3915763 100644 --- a/pytensor_distributions/vonmises.py +++ b/pytensor_distributions/vonmises.py @@ -68,7 +68,7 @@ def logpdf(x, mu, kappa): def rvs(mu, kappa, size=None, random_state=None): - return pt.random.vonmises(mu, kappa, rng=random_state, size=size) + return pt.random.vonmises(mu, kappa, rng=random_state, size=size, return_next_rng=True)[1] def cdf(x, mu, kappa): diff --git a/pytensor_distributions/wald.py b/pytensor_distributions/wald.py index fc15023..e4b9a9f 100644 --- a/pytensor_distributions/wald.py +++ b/pytensor_distributions/wald.py @@ -85,7 +85,7 @@ def sf(x, mu, lam): def rvs(mu, lam, size=None, random_state=None): - return pt.random.wald(mu, lam, rng=random_state, size=size) + return pt.random.wald(mu, lam, rng=random_state, size=size, return_next_rng=True)[1] def logcdf(x, mu, lam): diff --git a/pytensor_distributions/weibull.py b/pytensor_distributions/weibull.py index ff8fb0e..393ffc6 100644 --- a/pytensor_distributions/weibull.py +++ b/pytensor_distributions/weibull.py @@ -88,7 +88,7 @@ def sf(x, alpha, beta): def rvs(alpha, beta, size=None, random_state=None): - return pt.random.weibull(alpha, rng=random_state, size=size) * beta + return pt.random.weibull(alpha, rng=random_state, size=size, return_next_rng=True)[1] * beta def logcdf(x, alpha, beta): diff --git a/pytensor_distributions/wishart.py b/pytensor_distributions/wishart.py index d94003a..f4c57d0 100644 --- a/pytensor_distributions/wishart.py +++ b/pytensor_distributions/wishart.py @@ -82,12 +82,21 @@ def rvs(nu, V, size=None, random_state=None): squeeze = False chi_samples = pt.stack( - [pt.sqrt(pt.random.chisquare(nu - i, size=batch_size, rng=random_state)) for i in range(p)], + [ + pt.sqrt( + pt.random.chisquare( + nu - i, size=batch_size, rng=random_state, return_next_rng=True + )[1] + ) + for i in range(p) + ], axis=1, ) n_tril = p * (p - 1) // 2 - norm_samples = pt.random.normal(0, 1, size=(batch_size, n_tril), rng=random_state) + norm_samples = pt.random.normal( + 0, 1, size=(batch_size, n_tril), rng=random_state, return_next_rng=True + )[1] A = pt.zeros((batch_size, p, p)) diag_idx = pt.arange(p) diff --git a/pytensor_distributions/zi_binomial.py b/pytensor_distributions/zi_binomial.py index 6861d0b..27595a9 100644 --- a/pytensor_distributions/zi_binomial.py +++ b/pytensor_distributions/zi_binomial.py @@ -127,8 +127,10 @@ def isf(q, psi, n, p): def rvs(psi, n, p, size=None, random_state=None): - base_samples = pt.random.binomial(n, p, size=size, rng=random_state) - mask = pt.random.bernoulli(psi, size=size) + next_rng, base_samples = pt.random.binomial( + n, p, size=size, rng=random_state, return_next_rng=True + ) + next_rng, mask = pt.random.bernoulli(psi, size=size, rng=next_rng, return_next_rng=True) return pt.cast(mask, "int64") * base_samples diff --git a/pytensor_distributions/zi_negativebinomial.py b/pytensor_distributions/zi_negativebinomial.py index f9a4bb1..28668e0 100644 --- a/pytensor_distributions/zi_negativebinomial.py +++ b/pytensor_distributions/zi_negativebinomial.py @@ -148,9 +148,11 @@ def isf(q, psi, n, p): def rvs(psi, n, p, size=None, random_state=None): - base_samples = pt.random.negative_binomial(n, p, size=size, rng=random_state) - mask = pt.random.bernoulli(psi, size=size) - return pt.cast(mask, "int64") * base_samples + next_rng, nbase_samples = pt.random.negative_binomial( + n, p, size=size, rng=random_state, return_next_rng=True + ) + next_rng, mask = pt.random.bernoulli(psi, size=size, rng=next_rng, return_next_rng=True) + return pt.cast(mask, "int64") * nbase_samples def logcdf(x, psi, n, p): diff --git a/pytensor_distributions/zi_poisson.py b/pytensor_distributions/zi_poisson.py index 229fb87..5f7d2fb 100644 --- a/pytensor_distributions/zi_poisson.py +++ b/pytensor_distributions/zi_poisson.py @@ -129,8 +129,10 @@ def isf(q, psi, mu): def rvs(psi, mu, size=None, random_state=None): - base_samples = pt.random.poisson(mu, size=size, rng=random_state) - mask = pt.random.bernoulli(psi, size=size) + next_rng, base_samples = pt.random.poisson( + mu, size=size, rng=random_state, return_next_rng=True + ) + next_rng, mask = pt.random.bernoulli(psi, size=size, rng=next_rng, return_next_rng=True) return pt.cast(mask, "int64") * base_samples diff --git a/tests/test_polyagamma.py b/tests/test_polyagamma.py index 1401333..79d8507 100644 --- a/tests/test_polyagamma.py +++ b/tests/test_polyagamma.py @@ -67,7 +67,6 @@ def test_polyagamma_cdf_bounds(params, samples): assert_allclose(PolyaGamma.cdf(extended_vals, *p_params).eval(), expected) -@pytest.mark.slow @pytest.mark.parametrize("params", [(1.0, 0.0)]) def test_polyagamma_ppf(params, samples): """PPF should match empirical quantiles. Slow due to scan-based solver.""" @@ -79,7 +78,6 @@ def test_polyagamma_ppf(params, samples): assert_allclose(theoretical, empirical, rtol=1e-1, atol=5e-3) -@pytest.mark.slow def test_polyagamma_ppf_cdf_inverse(samples): """CDF(PPF(q)) should recover q. Slow due to scan-based solver.""" p_params, _ = samples[(1.0, 0.0)] diff --git a/tests/test_wishart.py b/tests/test_wishart.py index ab06b00..8faeb6b 100644 --- a/tests/test_wishart.py +++ b/tests/test_wishart.py @@ -120,9 +120,11 @@ def test_wishart_rvs(nu, V): n_samples = 1000 samples = Wishart.rvs(p_nu, p_V, size=n_samples).eval() - assert samples.shape == (n_samples, p, p), ( - f"Multiple samples should have shape ({n_samples}, {p}, {p})" - ) + assert samples.shape == ( + n_samples, + p, + p, + ), f"Multiple samples should have shape ({n_samples}, {p}, {p})" for i in range(min(10, n_samples)): eigenvalues = np.linalg.eigvalsh(samples[i])