diff --git a/docs/source/contributing/implementing_distribution.md b/docs/source/contributing/implementing_distribution.md index 85347622cdd..b0ca0c227b7 100644 --- a/docs/source/contributing/implementing_distribution.md +++ b/docs/source/contributing/implementing_distribution.md @@ -146,12 +146,12 @@ class Blah(PositiveContinuous): # We pass the standard parametrizations to super().dist @classmethod def dist(cls, param1, param2=None, alt_param2=None, **kwargs): - param1 = pt.as_tensor_variable(intX(param1)) + param1 = intX(param1) if param2 is not None and alt_param2 is not None: raise ValueError("Only one of param2 and alt_param2 is allowed.") if alt_param2 is not None: param2 = 1 / alt_param2 - param2 = pt.as_tensor_variable(floatX(param2)) + param2 = floatX(param2) # The first value-only argument should be a list of the parameters that # the rv_op needs in order to be instantiated diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py index 8688cb04c2f..cc1cb385b45 100644 --- a/pymc/distributions/continuous.py +++ b/pymc/distributions/continuous.py @@ -304,8 +304,8 @@ class Uniform(BoundedContinuous): @classmethod def dist(cls, lower=0, upper=1, **kwargs): - lower = pt.as_tensor_variable(floatX(lower)) - upper = pt.as_tensor_variable(floatX(upper)) + lower = floatX(lower) + upper = floatX(upper) return super().dist([lower, upper], **kwargs) def moment(rv, size, lower, upper): @@ -509,12 +509,8 @@ class Normal(Continuous): @classmethod def dist(cls, mu=0, sigma=None, tau=None, **kwargs): tau, sigma = get_tau_sigma(tau=tau, sigma=sigma) - sigma = pt.as_tensor_variable(sigma) - - # tau = pt.as_tensor_variable(tau) - # mean = median = mode = mu = pt.as_tensor_variable(floatX(mu)) - # variance = 1.0 / self.tau - + mu = floatX(mu) + sigma = floatX(sigma) return super().dist([mu, sigma], **kwargs) def moment(rv, size, mu, sigma): @@ -670,12 +666,11 @@ def dist( **kwargs, ) -> RandomVariable: tau, sigma = get_tau_sigma(tau=tau, sigma=sigma) - sigma = pt.as_tensor_variable(sigma) - tau = pt.as_tensor_variable(tau) - mu = pt.as_tensor_variable(floatX(mu)) + sigma = floatX(sigma) + mu = floatX(mu) - lower = pt.as_tensor_variable(floatX(lower)) if lower is not None else pt.constant(-np.inf) - upper = pt.as_tensor_variable(floatX(upper)) if upper is not None else pt.constant(np.inf) + lower = floatX(lower) if lower is not None else pt.constant(-np.inf) + upper = floatX(upper) if upper is not None else pt.constant(np.inf) return super().dist([mu, sigma, lower, upper], **kwargs) def moment(rv, size, mu, sigma, lower, upper): @@ -948,9 +943,9 @@ class Wald(PositiveContinuous): @classmethod def dist(cls, mu=None, lam=None, phi=None, alpha=0.0, **kwargs): mu, lam, phi = cls.get_mu_lam_phi(mu, lam, phi) - alpha = pt.as_tensor_variable(floatX(alpha)) - mu = pt.as_tensor_variable(floatX(mu)) - lam = pt.as_tensor_variable(floatX(lam)) + alpha = floatX(alpha) + mu = floatX(mu) + lam = floatX(lam) return super().dist([mu, lam, alpha], **kwargs) def moment(rv, size, mu, lam, alpha): @@ -1115,8 +1110,8 @@ class Beta(UnitContinuous): @classmethod def dist(cls, alpha=None, beta=None, mu=None, sigma=None, nu=None, *args, **kwargs): alpha, beta = cls.get_alpha_beta(alpha, beta, mu, sigma, nu) - alpha = pt.as_tensor_variable(floatX(alpha)) - beta = pt.as_tensor_variable(floatX(beta)) + alpha = floatX(alpha) + beta = floatX(beta) return super().dist([alpha, beta], **kwargs) @@ -1242,8 +1237,8 @@ class Kumaraswamy(UnitContinuous): @classmethod def dist(cls, a, b, *args, **kwargs): - a = pt.as_tensor_variable(floatX(a)) - b = pt.as_tensor_variable(floatX(b)) + a = floatX(a) + b = floatX(b) return super().dist([a, b], *args, **kwargs) @@ -1328,7 +1323,7 @@ class Exponential(PositiveContinuous): @classmethod def dist(cls, lam, *args, **kwargs): - lam = pt.as_tensor_variable(floatX(lam)) + lam = floatX(lam) # PyTensor exponential op is parametrized in terms of mu (1/lam) return super().dist([pt.reciprocal(lam)], **kwargs) @@ -1409,8 +1404,8 @@ class Laplace(Continuous): @classmethod def dist(cls, mu, b, *args, **kwargs): - b = pt.as_tensor_variable(floatX(b)) - mu = pt.as_tensor_variable(floatX(mu)) + b = floatX(b) + mu = floatX(mu) return super().dist([mu, b], *args, **kwargs) @@ -1518,9 +1513,9 @@ class AsymmetricLaplace(Continuous): @classmethod def dist(cls, kappa=None, mu=None, b=None, q=None, *args, **kwargs): kappa = cls.get_kappa(kappa, q) - b = pt.as_tensor_variable(floatX(b)) - kappa = pt.as_tensor_variable(floatX(kappa)) - mu = pt.as_tensor_variable(floatX(mu)) + b = floatX(b) + kappa = floatX(kappa) + mu = floatX(mu) return super().dist([b, kappa, mu], *args, **kwargs) @@ -1634,8 +1629,8 @@ class LogNormal(PositiveContinuous): def dist(cls, mu=0, sigma=None, tau=None, *args, **kwargs): tau, sigma = get_tau_sigma(tau=tau, sigma=sigma) - mu = pt.as_tensor_variable(floatX(mu)) - sigma = pt.as_tensor_variable(floatX(sigma)) + mu = floatX(mu) + sigma = floatX(sigma) return super().dist([mu, sigma], *args, **kwargs) @@ -1759,9 +1754,9 @@ class StudentT(Continuous): @classmethod def dist(cls, nu, mu=0, *, sigma=None, lam=None, **kwargs): - nu = pt.as_tensor_variable(floatX(nu)) + nu = floatX(nu) lam, sigma = get_tau_sigma(tau=lam, sigma=sigma) - sigma = pt.as_tensor_variable(sigma) + sigma = floatX(sigma) return super().dist([nu, mu, sigma], **kwargs) @@ -1856,8 +1851,8 @@ class Pareto(BoundedContinuous): @classmethod def dist(cls, alpha, m, **kwargs): - alpha = pt.as_tensor_variable(floatX(alpha)) - m = pt.as_tensor_variable(floatX(m)) + alpha = floatX(alpha) + m = floatX(m) return super().dist([alpha, m], **kwargs) @@ -1953,8 +1948,8 @@ class Cauchy(Continuous): @classmethod def dist(cls, alpha, beta, *args, **kwargs): - alpha = pt.as_tensor_variable(floatX(alpha)) - beta = pt.as_tensor_variable(floatX(beta)) + alpha = floatX(alpha) + beta = floatX(beta) return super().dist([alpha, beta], **kwargs) @@ -2024,7 +2019,7 @@ class HalfCauchy(PositiveContinuous): @classmethod def dist(cls, beta, *args, **kwargs): - beta = pt.as_tensor_variable(floatX(beta)) + beta = floatX(beta) return super().dist([0.0, beta], **kwargs) def moment(rv, size, loc, beta): @@ -2119,8 +2114,8 @@ class Gamma(PositiveContinuous): @classmethod def dist(cls, alpha=None, beta=None, mu=None, sigma=None, **kwargs): alpha, beta = cls.get_alpha_beta(alpha, beta, mu, sigma) - alpha = pt.as_tensor_variable(floatX(alpha)) - beta = pt.as_tensor_variable(floatX(beta)) + alpha = floatX(alpha) + beta = floatX(beta) # The PyTensor `GammaRV` `Op` will invert the `beta` parameter itself return super().dist([alpha, beta], **kwargs) @@ -2228,8 +2223,8 @@ class InverseGamma(PositiveContinuous): @classmethod def dist(cls, alpha=None, beta=None, mu=None, sigma=None, *args, **kwargs): alpha, beta = cls._get_alpha_beta(alpha, beta, mu, sigma) - alpha = pt.as_tensor_variable(floatX(alpha)) - beta = pt.as_tensor_variable(floatX(beta)) + alpha = floatX(alpha) + beta = floatX(beta) return super().dist([alpha, beta], **kwargs) @@ -2332,7 +2327,7 @@ class ChiSquared(PositiveContinuous): @classmethod def dist(cls, nu, *args, **kwargs): - nu = pt.as_tensor_variable(floatX(nu)) + nu = floatX(nu) return super().dist([nu], *args, **kwargs) def moment(rv, size, nu): @@ -2417,8 +2412,8 @@ class Weibull(PositiveContinuous): @classmethod def dist(cls, alpha, beta, *args, **kwargs): - alpha = pt.as_tensor_variable(floatX(alpha)) - beta = pt.as_tensor_variable(floatX(beta)) + alpha = floatX(alpha) + beta = floatX(beta) return super().dist([alpha, beta], *args, **kwargs) @@ -2537,9 +2532,9 @@ class HalfStudentT(PositiveContinuous): @classmethod def dist(cls, nu, sigma=None, lam=None, *args, **kwargs): - nu = pt.as_tensor_variable(floatX(nu)) + nu = floatX(nu) lam, sigma = get_tau_sigma(lam, sigma) - sigma = pt.as_tensor_variable(sigma) + sigma = floatX(sigma) return super().dist([nu, sigma], *args, **kwargs) @@ -2657,9 +2652,9 @@ class ExGaussian(Continuous): @classmethod def dist(cls, mu=0.0, sigma=None, nu=None, *args, **kwargs): - mu = pt.as_tensor_variable(floatX(mu)) - sigma = pt.as_tensor_variable(floatX(sigma)) - nu = pt.as_tensor_variable(floatX(nu)) + mu = floatX(mu) + sigma = floatX(sigma) + nu = floatX(nu) return super().dist([mu, sigma, nu], *args, **kwargs) @@ -2762,8 +2757,8 @@ class VonMises(CircularContinuous): @classmethod def dist(cls, mu=0.0, kappa=1.0, *args, **kwargs): - mu = pt.as_tensor_variable(floatX(mu)) - kappa = pt.as_tensor_variable(floatX(kappa)) + mu = floatX(mu) + kappa = floatX(kappa) return super().dist([mu, kappa], *args, **kwargs) def moment(rv, size, mu, kappa): @@ -2865,10 +2860,9 @@ class SkewNormal(Continuous): @classmethod def dist(cls, alpha=1, mu=0.0, sigma=None, tau=None, *args, **kwargs): tau, sigma = get_tau_sigma(tau=tau, sigma=sigma) - alpha = pt.as_tensor_variable(floatX(alpha)) - mu = pt.as_tensor_variable(floatX(mu)) - tau = pt.as_tensor_variable(tau) - sigma = pt.as_tensor_variable(sigma) + alpha = floatX(alpha) + mu = floatX(mu) + sigma = floatX(sigma) return super().dist([mu, sigma, alpha], *args, **kwargs) @@ -2954,9 +2948,9 @@ class Triangular(BoundedContinuous): @classmethod def dist(cls, lower=0, upper=1, c=0.5, *args, **kwargs): - lower = pt.as_tensor_variable(floatX(lower)) - upper = pt.as_tensor_variable(floatX(upper)) - c = pt.as_tensor_variable(floatX(c)) + lower = floatX(lower) + upper = floatX(upper) + c = floatX(c) return super().dist([lower, c, upper], *args, **kwargs) @@ -3061,8 +3055,8 @@ class Gumbel(Continuous): @classmethod def dist(cls, mu, beta, **kwargs): - mu = pt.as_tensor_variable(floatX(mu)) - beta = pt.as_tensor_variable(floatX(beta)) + mu = floatX(mu) + beta = floatX(beta) return super().dist([mu, beta], **kwargs) @@ -3171,8 +3165,8 @@ class Rice(PositiveContinuous): @classmethod def dist(cls, nu=None, sigma=None, b=None, *args, **kwargs): nu, b, sigma = cls.get_nu_b(nu, b, sigma) - b = pt.as_tensor_variable(floatX(b)) - sigma = pt.as_tensor_variable(floatX(sigma)) + b = floatX(b) + sigma = floatX(sigma) return super().dist([b, sigma], *args, **kwargs) @@ -3270,8 +3264,8 @@ class Logistic(Continuous): @classmethod def dist(cls, mu=0.0, s=1.0, *args, **kwargs): - mu = pt.as_tensor_variable(floatX(mu)) - s = pt.as_tensor_variable(floatX(s)) + mu = floatX(mu) + s = floatX(s) return super().dist([mu, s], *args, **kwargs) def moment(rv, size, mu, s): @@ -3365,11 +3359,9 @@ class LogitNormal(UnitContinuous): @classmethod def dist(cls, mu=0, sigma=None, tau=None, **kwargs): - mu = pt.as_tensor_variable(floatX(mu)) + mu = floatX(mu) tau, sigma = get_tau_sigma(tau=tau, sigma=sigma) - sigma = pt.as_tensor_variable(sigma) - tau = pt.as_tensor_variable(tau) - + sigma = floatX(sigma) return super().dist([mu, sigma], **kwargs) def moment(rv, size, mu, sigma): @@ -3499,11 +3491,6 @@ def dist(cls, x_points, pdf_points, *args, **kwargs): x_points = pt.constant(floatX(x_points)) pdf_points = pt.constant(floatX(pdf_points)) cdf_points = pt.constant(floatX(cdf_points)) - - # lower = pt.as_tensor_variable(x_points[0]) - # upper = pt.as_tensor_variable(x_points[-1]) - # median = _interpolated_argcdf(0.5, pdf_points, cdf_points, x_points) - return super().dist([x_points, pdf_points, cdf_points], **kwargs) def moment(rv, size, x_points, pdf_points, cdf_points): @@ -3607,8 +3594,8 @@ class Moyal(Continuous): @classmethod def dist(cls, mu=0, sigma=1.0, *args, **kwargs): - mu = pt.as_tensor_variable(floatX(mu)) - sigma = pt.as_tensor_variable(floatX(sigma)) + mu = floatX(mu) + sigma = floatX(sigma) return super().dist([mu, sigma], *args, **kwargs) @@ -3694,9 +3681,9 @@ def __init__(self, get_pdf=False): self.get_pdf = get_pdf def make_node(self, x, h, z): - x = pt.as_tensor_variable(floatX(x)) - h = pt.as_tensor_variable(floatX(h)) - z = pt.as_tensor_variable(floatX(z)) + x = floatX(x) + h = floatX(h) + z = floatX(z) bshape = broadcast_shape(x, h, z) shape = [None] * len(bshape) return Apply(self, [x, h, z], [pt.TensorType(pytensor.config.floatX, shape)()]) @@ -3795,8 +3782,8 @@ class PolyaGamma(PositiveContinuous): @classmethod def dist(cls, h=1.0, z=0.0, **kwargs): - h = pt.as_tensor_variable(floatX(h)) - z = pt.as_tensor_variable(floatX(z)) + h = floatX(h) + z = floatX(z) msg = f"The variable {h} specified for PolyaGamma has non-positive " msg += "values, making it unsuitable for this parameter." diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index 8e52c812d16..8f59ff15bfe 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -131,8 +131,8 @@ def dist(cls, n, p=None, logit_p=None, *args, **kwargs): if logit_p is not None: p = pt.sigmoid(logit_p) - n = pt.as_tensor_variable(intX(n)) - p = pt.as_tensor_variable(floatX(p)) + n = intX(n) + p = floatX(p) return super().dist([n, p], **kwargs) def moment(rv, size, n, p): @@ -239,9 +239,9 @@ def BetaBinom(a, b, n, x): @classmethod def dist(cls, alpha, beta, n, *args, **kwargs): - alpha = pt.as_tensor_variable(floatX(alpha)) - beta = pt.as_tensor_variable(floatX(beta)) - n = pt.as_tensor_variable(intX(n)) + alpha = floatX(alpha) + beta = floatX(beta) + n = intX(n) return super().dist([n, alpha, beta], **kwargs) def moment(rv, size, n, alpha, beta): @@ -353,7 +353,7 @@ def dist(cls, p=None, logit_p=None, *args, **kwargs): if logit_p is not None: p = pt.sigmoid(logit_p) - p = pt.as_tensor_variable(floatX(p)) + p = floatX(p) return super().dist([p], **kwargs) def moment(rv, size, p): @@ -461,8 +461,8 @@ def DiscreteWeibull(q, b, x): @classmethod def dist(cls, q, beta, *args, **kwargs): - q = pt.as_tensor_variable(floatX(q)) - beta = pt.as_tensor_variable(floatX(beta)) + q = floatX(q) + beta = floatX(beta) return super().dist([q, beta], **kwargs) def moment(rv, size, q, beta): @@ -550,7 +550,7 @@ class Poisson(Discrete): @classmethod def dist(cls, mu, *args, **kwargs): - mu = pt.as_tensor_variable(floatX(mu)) + mu = floatX(mu) return super().dist([mu], *args, **kwargs) def moment(rv, size, mu): @@ -673,8 +673,8 @@ def NegBinom(a, m, x): @classmethod def dist(cls, mu=None, alpha=None, p=None, n=None, *args, **kwargs): n, p = cls.get_n_p(mu=mu, alpha=alpha, p=p, n=n) - n = pt.as_tensor_variable(floatX(n)) - p = pt.as_tensor_variable(floatX(p)) + n = floatX(n) + p = floatX(p) return super().dist([n, p], *args, **kwargs) @classmethod @@ -785,7 +785,7 @@ class Geometric(Discrete): @classmethod def dist(cls, p, *args, **kwargs): - p = pt.as_tensor_variable(floatX(p)) + p = floatX(p) return super().dist([p], *args, **kwargs) def moment(rv, size, p): @@ -882,9 +882,9 @@ class HyperGeometric(Discrete): @classmethod def dist(cls, N, k, n, *args, **kwargs): - good = pt.as_tensor_variable(intX(k)) - bad = pt.as_tensor_variable(intX(N - k)) - n = pt.as_tensor_variable(intX(n)) + good = intX(k) + bad = intX(N - k) + n = intX(n) return super().dist([good, bad, n], *args, **kwargs) def moment(rv, size, good, bad, n): @@ -1117,7 +1117,7 @@ def dist(cls, p=None, logit_p=None, **kwargs): if logit_p is not None: p = pm.math.softmax(logit_p, axis=-1) - p = pt.as_tensor_variable(p) + p = floatX(p) if isinstance(p, TensorConstant): p_ = np.asarray(p.data) if np.any(p_ < 0): @@ -1131,7 +1131,7 @@ def dist(cls, p=None, logit_p=None, **kwargs): UserWarning, ) p_ = p_ / pt.sum(p_, axis=-1, keepdims=True) - p = pt.as_tensor_variable(p_) + p = floatX(p_) return super().dist([p], **kwargs) def moment(rv, size, p): @@ -1258,7 +1258,7 @@ def _zero_inflated_mixture(*, name, nonzero_p, nonzero_dist, **kwargs): If name is `None`, this function returns an unregistered variable """ - nonzero_p = pt.as_tensor_variable(floatX(nonzero_p)) + nonzero_p = floatX(nonzero_p) weights = pt.stack([1 - nonzero_p, nonzero_p], axis=-1) comp_dists = [ DiracDelta.dist(0), @@ -1508,8 +1508,8 @@ class _OrderedLogistic(Categorical): @classmethod def dist(cls, eta, cutpoints, *args, **kwargs): - eta = pt.as_tensor_variable(floatX(eta)) - cutpoints = pt.as_tensor_variable(cutpoints) + eta = floatX(eta) + cutpoints = floatX(cutpoints) pa = sigmoid(cutpoints - pt.shape_padright(eta)) p_cum = pt.concatenate( @@ -1614,8 +1614,8 @@ class _OrderedProbit(Categorical): @classmethod def dist(cls, eta, cutpoints, sigma=1, *args, **kwargs): - eta = pt.as_tensor_variable(floatX(eta)) - cutpoints = pt.as_tensor_variable(cutpoints) + eta = floatX(eta) + cutpoints = floatX(cutpoints) probits = pt.shape_padright(eta) - cutpoints _log_p = pt.concatenate( @@ -1628,7 +1628,7 @@ def dist(cls, eta, cutpoints, sigma=1, *args, **kwargs): ], axis=-1, ) - _log_p = pt.as_tensor_variable(floatX(_log_p)) + _log_p = floatX(_log_p) p = pt.exp(_log_p) return super().dist(p, *args, **kwargs) diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 7f873a2e4a1..d0ef3abb5ac 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -258,7 +258,7 @@ class MvNormal(Continuous): @classmethod def dist(cls, mu, cov=None, tau=None, chol=None, lower=True, **kwargs): - mu = pt.as_tensor_variable(mu) + mu = floatX(mu) cov = quaddist_matrix(cov, chol, tau, lower) # PyTensor is stricter about the shape of mu, than PyMC used to be mu = pt.broadcast_arrays(mu, cov[..., -1])[0] @@ -395,8 +395,8 @@ def dist(cls, nu, Sigma=None, mu=None, scale=None, tau=None, chol=None, lower=Tr if scale is not None: raise ValueError("Specify only one of scale and Sigma") scale = Sigma - nu = pt.as_tensor_variable(floatX(nu)) - mu = pt.as_tensor_variable(floatX(mu)) + nu = floatX(nu) + mu = floatX(mu) scale = quaddist_matrix(scale, chol, tau, lower) # PyTensor is stricter about the shape of mu, than PyMC used to be mu = pt.broadcast_arrays(mu, scale[..., -1])[0] @@ -462,10 +462,7 @@ class Dirichlet(SimplexContinuous): @classmethod def dist(cls, a, **kwargs): - a = pt.as_tensor_variable(a) - # mean = a / pt.sum(a) - # mode = pt.switch(pt.all(a > 1), (a - 1) / pt.sum(a - 1), np.nan) - + a = floatX(a) return super().dist([a], **kwargs) def moment(rv, size, a): @@ -541,7 +538,7 @@ class Multinomial(Discrete): @classmethod def dist(cls, n, p, *args, **kwargs): - p = pt.as_tensor_variable(p) + p = floatX(p) if isinstance(p, TensorConstant): p_ = np.asarray(p.data) if np.any(p_ < 0): @@ -555,9 +552,8 @@ def dist(cls, n, p, *args, **kwargs): UserWarning, ) p_ = p_ / pt.sum(p_, axis=-1, keepdims=True) - p = pt.as_tensor_variable(p_) - n = pt.as_tensor_variable(n) - p = pt.as_tensor_variable(p) + p = floatX(p_) + n = intX(n) return super().dist([n, p], *args, **kwargs) def moment(rv, size, n, p): @@ -729,9 +725,9 @@ class _OrderedMultinomial(Multinomial): @classmethod def dist(cls, eta, cutpoints, n, *args, **kwargs): - eta = pt.as_tensor_variable(floatX(eta)) - cutpoints = pt.as_tensor_variable(cutpoints) - n = pt.as_tensor_variable(intX(n)) + eta = floatX(eta) + cutpoints = floatX(cutpoints) + n = intX(n) pa = sigmoid(cutpoints - pt.shape_padright(eta)) p_cum = pt.concatenate( @@ -948,8 +944,8 @@ class Wishart(Continuous): @classmethod def dist(cls, nu, V, *args, **kwargs): - nu = pt.as_tensor_variable(intX(nu)) - V = pt.as_tensor_variable(floatX(V)) + nu = intX(nu) + V = floatX(V) warnings.warn( "The Wishart distribution can currently not be used " @@ -1179,8 +1175,8 @@ class _LKJCholeskyCov(Distribution): @classmethod def dist(cls, n, eta, sd_dist, **kwargs): - n = pt.as_tensor_variable(intX(n)) - eta = pt.as_tensor_variable(floatX(eta)) + n = intX(n) + eta = floatX(eta) if not ( isinstance(sd_dist, Variable) @@ -1577,8 +1573,8 @@ class LKJCorr(BoundedContinuous): @classmethod def dist(cls, n, eta, **kwargs): - n = pt.as_tensor_variable(intX(n)) - eta = pt.as_tensor_variable(floatX(eta)) + n = intX(n) + eta = floatX(eta) return super().dist([n, eta], **kwargs) def moment(rv, *args): @@ -1783,7 +1779,7 @@ def dist( else: if rowchol.ndim != 2: raise ValueError("rowchol must be two dimensional.") - rowchol_cov = pt.as_tensor_variable(rowchol) + rowchol_cov = floatX(rowchol) # Among-column matrices if len([i for i in [colcov, colchol] if i is not None]) != 1: @@ -1791,20 +1787,20 @@ def dist( "Incompatible parameterization. Specify exactly one of colcov, or colchol." ) if colcov is not None: - colcov = pt.as_tensor_variable(colcov) + colcov = floatX(colcov) if colcov.ndim != 2: raise ValueError("colcov must be two dimensional.") colchol_cov = cholesky(colcov) else: if colchol.ndim != 2: raise ValueError("colchol must be two dimensional.") - colchol_cov = pt.as_tensor_variable(colchol) + colchol_cov = floatX(colchol) dist_shape = (rowchol_cov.shape[-1], colchol_cov.shape[-1]) # Broadcasting mu mu = pt.extra_ops.broadcast_to(mu, shape=dist_shape) - mu = pt.as_tensor_variable(floatX(mu)) + mu = floatX(mu) return super().dist([mu, rowchol_cov, colchol_cov], **kwargs) @@ -1983,7 +1979,7 @@ def dist(cls, mu, covs=None, chols=None, evds=None, sigma=None, *args, **kwargs) cov_i = pt.dot(Q, pt.dot(pt.diag(eig), Q.T)) covs.append(cov_i) - mu = pt.as_tensor_variable(mu) + mu = floatX(mu) return super().dist([mu, sigma, *covs], **kwargs) @@ -2050,7 +2046,7 @@ class CARRV(RandomVariable): _print_name = ("CAR", "\\operatorname{CAR}") def make_node(self, rng, size, dtype, mu, W, alpha, tau): - mu = pt.as_tensor_variable(floatX(mu)) + mu = floatX(mu) W = pytensor.sparse.as_sparse_or_tensor_variable(floatX(W)) if not W.ndim == 2: @@ -2064,8 +2060,8 @@ def make_node(self, rng, size, dtype, mu, W, alpha, tau): else: W = Assert(msg)(W, pt.allclose(W, W.T)) - tau = pt.as_tensor_variable(floatX(tau)) - alpha = pt.as_tensor_variable(floatX(alpha)) + tau = floatX(tau) + alpha = floatX(alpha) return super().make_node(rng, size, dtype, mu, W, alpha, tau) @@ -2231,7 +2227,7 @@ class StickBreakingWeightsRV(RandomVariable): def make_node(self, rng, size, dtype, alpha, K): alpha = pt.as_tensor_variable(alpha) - K = pt.as_tensor_variable(intX(K)) + K = intX(K) if K.ndim > 0: raise ValueError("K must be a scalar.") @@ -2312,8 +2308,8 @@ class StickBreakingWeights(SimplexContinuous): @classmethod def dist(cls, alpha, K, *args, **kwargs): - alpha = pt.as_tensor_variable(floatX(alpha)) - K = pt.as_tensor_variable(intX(K)) + alpha = floatX(alpha) + K = intX(K) return super().dist([alpha, K], **kwargs) @@ -2488,7 +2484,7 @@ def __new__( def dist(cls, sigma=1, n_zerosum_axes=None, support_shape=None, **kwargs): n_zerosum_axes = cls.check_zerosum_axes(n_zerosum_axes) - sigma = pt.as_tensor_variable(floatX(sigma)) + sigma = floatX(sigma) if sigma.ndim > 0: raise ValueError("sigma has to be a scalar") @@ -2504,7 +2500,7 @@ def dist(cls, sigma=1, n_zerosum_axes=None, support_shape=None, **kwargs): # TODO: edge-case doesn't work for now, because pt.stack in get_support_shape fails # else: # support_shape = () # because it's just a Normal in that case - support_shape = pt.as_tensor_variable(intX(support_shape)) + support_shape = intX(support_shape) assert n_zerosum_axes == pt.get_vector_length( support_shape diff --git a/pymc/distributions/simulator.py b/pymc/distributions/simulator.py index 4ae0c9fb030..7b806be3d61 100644 --- a/pymc/distributions/simulator.py +++ b/pymc/distributions/simulator.py @@ -199,7 +199,7 @@ def dist( # type: ignore else: raise ValueError(f"The summary statistic {sum_stat} is not implemented") - epsilon = pt.as_tensor_variable(floatX(epsilon)) + epsilon = floatX(epsilon) if params is None: params = unnamed_params diff --git a/pymc/distributions/timeseries.py b/pymc/distributions/timeseries.py index 1cee128881b..fcd38b17eab 100644 --- a/pymc/distributions/timeseries.py +++ b/pymc/distributions/timeseries.py @@ -125,7 +125,7 @@ def dist(cls, init_dist, innovation_dist, steps=None, **kwargs) -> pt.TensorVari ) if steps is None: raise ValueError("Must specify steps or shape parameter") - steps = pt.as_tensor_variable(intX(steps)) + steps = intX(steps) return super().dist([init_dist, innovation_dist, steps], **kwargs) @@ -305,8 +305,8 @@ def get_dists(cls, mu=0.0, sigma=1.0, *, init_dist=None, **kwargs): ) init_dist = Normal.dist(0, 100) - mu = pt.as_tensor_variable(mu) - sigma = pt.as_tensor_variable(sigma) + mu = floatX(mu) + sigma = floatX(sigma) innovation_dist = Normal.dist(mu=mu, sigma=sigma) return init_dist, innovation_dist, kwargs @@ -506,7 +506,7 @@ class AR(Distribution): rv_type = AutoRegressiveRV def __new__(cls, name, rho, *args, steps=None, constant=False, ar_order=None, **kwargs): - rhos = pt.atleast_1d(pt.as_tensor_variable(floatX(rho))) + rhos = pt.atleast_1d(floatX(rho)) ar_order = cls._get_ar_order(rhos=rhos, constant=constant, ar_order=ar_order) steps = get_support_shape_1d( support_shape=steps, @@ -533,8 +533,8 @@ def dist( **kwargs, ): _, sigma = get_tau_sigma(tau=tau, sigma=sigma) - sigma = pt.as_tensor_variable(floatX(sigma)) - rhos = pt.atleast_1d(pt.as_tensor_variable(floatX(rho))) + sigma = floatX(sigma) + rhos = pt.atleast_1d(floatX(rho)) if "init" in kwargs: warnings.warn( @@ -786,10 +786,10 @@ def dist(cls, omega, alpha_1, beta_1, initial_vol, *, steps=None, **kwargs): raise ValueError("Must specify steps or shape parameter") steps = pt.as_tensor_variable(intX(steps), ndim=0) - omega = pt.as_tensor_variable(omega) - alpha_1 = pt.as_tensor_variable(alpha_1) - beta_1 = pt.as_tensor_variable(beta_1) - initial_vol = pt.as_tensor_variable(initial_vol) + omega = floatX(omega) + alpha_1 = floatX(alpha_1) + beta_1 = floatX(beta_1) + initial_vol = floatX(initial_vol) init_dist = Normal.dist(0, initial_vol) # We can ignore init_dist, as it will be accounted for in the logp term @@ -933,7 +933,7 @@ class EulerMaruyama(Distribution): rv_type = EulerMaruyamaRV def __new__(cls, name, dt, sde_fn, *args, steps=None, **kwargs): - dt = pt.as_tensor_variable(floatX(dt)) + dt = floatX(dt) steps = get_support_shape_1d( support_shape=steps, shape=None, # Shape will be checked in `cls.dist` @@ -952,8 +952,8 @@ def dist(cls, dt, sde_fn, sde_pars, *, init_dist=None, steps=None, **kwargs): raise ValueError("Must specify steps or shape parameter") steps = pt.as_tensor_variable(intX(steps), ndim=0) - dt = pt.as_tensor_variable(floatX(dt)) - sde_pars = [pt.as_tensor_variable(x) for x in sde_pars] + dt = floatX(dt) + sde_pars = [floatX(x) for x in sde_pars] if init_dist is not None: if not isinstance(init_dist, TensorVariable) or not isinstance( diff --git a/pymc/pytensorf.py b/pymc/pytensorf.py index 3c9cb945be9..d52e45eb389 100644 --- a/pymc/pytensorf.py +++ b/pymc/pytensorf.py @@ -435,8 +435,7 @@ def floatX(X): try: return X.astype(pytensor.config.floatX) except AttributeError: - # Scalar passed - return np.asarray(X, dtype=pytensor.config.floatX) + return pt.as_tensor_variable(X, dtype=pytensor.config.floatX) _conversion_map = {"float64": "int32", "float32": "int16", "float16": "int8", "float8": "int8"} @@ -450,8 +449,7 @@ def intX(X): try: return X.astype(intX) except AttributeError: - # Scalar passed - return np.asarray(X, dtype=intX) + return pt.as_tensor_variable(X, dtype=intX) def smartfloatX(x): diff --git a/tests/distributions/test_continuous.py b/tests/distributions/test_continuous.py index 8857a569fa2..f83b271bdc4 100644 --- a/tests/distributions/test_continuous.py +++ b/tests/distributions/test_continuous.py @@ -2284,3 +2284,10 @@ def dist(cls, **kwargs): extra_args={"rng": pytensor.shared(rng)}, ref_rand=ref_rand, ) + + +def test_list_of_variables_input(): + x = pm.HalfNormal.dist(1) + y = pm.HalfCauchy.dist(1) + z = pm.Exponential.dist(lam=[x, y]) + assert z.eval().shape == (2,)