From 3a304d6fcea326129824ce69b665b48bd8681d42 Mon Sep 17 00:00:00 2001 From: Anatoly <44327258+aerubanov@users.noreply.github.com> Date: Sat, 17 Feb 2024 10:30:33 +0400 Subject: [PATCH] Remove intX and floatX from distributions (#7114) --- pymc/distributions/bound.py | 18 ++-- pymc/distributions/continuous.py | 125 +++++++++++++------------- pymc/distributions/discrete.py | 41 +++++---- pymc/distributions/mixture.py | 5 +- pymc/distributions/multivariate.py | 59 ++++++------ pymc/distributions/simulator.py | 3 +- pymc/distributions/timeseries.py | 12 +-- pymc/distributions/transforms.py | 6 +- tests/distributions/test_mixture.py | 4 +- tests/distributions/test_truncated.py | 5 +- tests/test_initial_point.py | 5 +- tests/variational/test_inference.py | 7 +- 12 files changed, 144 insertions(+), 146 deletions(-) diff --git a/pymc/distributions/bound.py b/pymc/distributions/bound.py index 0f4379d97fb..ad8ec61f853 100644 --- a/pymc/distributions/bound.py +++ b/pymc/distributions/bound.py @@ -27,7 +27,6 @@ from pymc.distributions.transforms import _default_transform from pymc.logprob.basic import logp from pymc.model import modelcontext -from pymc.pytensorf import floatX, intX from pymc.util import check_dist_not_registered __all__ = ["Bound"] @@ -206,7 +205,7 @@ def __new__( res = _ContinuousBounded( name, [dist, lower, upper], - initval=floatX(initval), + initval=initval.astype("float"), size=size, shape=shape, **kwargs, @@ -215,7 +214,7 @@ def __new__( res = _DiscreteBounded( name, [dist, lower, upper], - initval=intX(initval), + initval=initval.astype("int"), size=size, shape=shape, **kwargs, @@ -241,7 +240,7 @@ def dist( shape=shape, **kwargs, ) - res.tag.test_value = floatX(initval) + res.tag.test_value = initval else: res = _DiscreteBounded.dist( [dist, lower, upper], @@ -249,7 +248,7 @@ def dist( shape=shape, **kwargs, ) - res.tag.test_value = intX(initval) + res.tag.test_value = initval.astype("int") return res @classmethod @@ -286,9 +285,9 @@ def _set_values(cls, lower, upper, size, shape, initval): size = shape lower = np.asarray(lower) - lower = floatX(np.where(lower == None, -np.inf, lower)) # noqa E711 + lower = np.where(lower == None, -np.inf, lower) # noqa E711 upper = np.asarray(upper) - upper = floatX(np.where(upper == None, np.inf, upper)) # noqa E711 + upper = np.where(upper == None, np.inf, upper) # noqa E711 if initval is None: _size = np.broadcast_shapes(to_tuple(size), np.shape(lower), np.shape(upper)) @@ -303,7 +302,6 @@ def _set_values(cls, lower, upper, size, shape, initval): np.where(_upper == np.inf, _lower + 1, (_lower + _upper) / 2), ), ) - - lower = as_tensor_variable(floatX(lower)) - upper = as_tensor_variable(floatX(upper)) + lower = as_tensor_variable(lower, dtype="floatX") + upper = as_tensor_variable(upper, dtype="floatX") return lower, upper, initval diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py index 5bb9b03eceb..19aac7f468e 100644 --- a/pymc/distributions/continuous.py +++ b/pymc/distributions/continuous.py @@ -96,7 +96,6 @@ def polyagamma_cdf(*args, **kwargs): from pymc.distributions.shape_utils import rv_size_is_none from pymc.distributions.transforms import _default_transform from pymc.math import invlogit, logdiffexp, logit -from pymc.pytensorf import floatX __all__ = [ "Uniform", @@ -257,7 +256,7 @@ def get_tau_sigma(tau=None, sigma=None): raise ValueError("tau must be positive") sigma = tau_**-0.5 - return floatX(tau), floatX(sigma) + return tau, sigma class Uniform(BoundedContinuous): @@ -309,8 +308,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 = pt.as_tensor_variable(lower) + upper = pt.as_tensor_variable(upper) return super().dist([lower, upper], **kwargs) def moment(rv, size, lower, upper): @@ -679,10 +678,10 @@ def dist( ) -> RandomVariable: tau, sigma = get_tau_sigma(tau=tau, sigma=sigma) sigma = pt.as_tensor_variable(sigma) - mu = pt.as_tensor_variable(floatX(mu)) + mu = pt.as_tensor_variable(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 = pt.as_tensor_variable(lower) if lower is not None else pt.constant(-np.inf) + upper = pt.as_tensor_variable(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): @@ -1000,9 +999,9 @@ def dist( **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 = pt.as_tensor_variable(alpha) + mu = pt.as_tensor_variable(mu) + lam = pt.as_tensor_variable(lam) return super().dist([mu, lam, alpha], **kwargs) def moment(rv, size, mu, lam, alpha): @@ -1176,8 +1175,8 @@ def dist( **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 = pt.as_tensor_variable(alpha) + beta = pt.as_tensor_variable(beta) return super().dist([alpha, beta], **kwargs) @@ -1304,8 +1303,8 @@ class Kumaraswamy(UnitContinuous): @classmethod def dist(cls, a: DIST_PARAMETER_TYPES, b: DIST_PARAMETER_TYPES, *args, **kwargs): - a = pt.as_tensor_variable(floatX(a)) - b = pt.as_tensor_variable(floatX(b)) + a = pt.as_tensor_variable(a) + b = pt.as_tensor_variable(b) return super().dist([a, b], *args, **kwargs) @@ -1401,7 +1400,7 @@ def dist(cls, lam=None, scale=None, *args, **kwargs): if scale is None: scale = pt.reciprocal(lam) - scale = pt.as_tensor_variable(floatX(scale)) + scale = pt.as_tensor_variable(scale) # PyTensor exponential op is parametrized in terms of mu (1/lam) return super().dist([scale], **kwargs) @@ -1491,8 +1490,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 = pt.as_tensor_variable(b) + mu = pt.as_tensor_variable(mu) return super().dist([mu, b], *args, **kwargs) @@ -1608,9 +1607,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 = pt.as_tensor_variable(b) + kappa = pt.as_tensor_variable(kappa) + mu = pt.as_tensor_variable(mu) return super().dist([b, kappa, mu], *args, **kwargs) @@ -1724,8 +1723,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 = pt.as_tensor_variable(mu) + sigma = pt.as_tensor_variable(sigma) return super().dist([mu, sigma], *args, **kwargs) @@ -1839,7 +1838,7 @@ class StudentT(Continuous): @classmethod def dist(cls, nu, mu=0, *, sigma=None, lam=None, **kwargs): - nu = pt.as_tensor_variable(floatX(nu)) + nu = pt.as_tensor_variable(nu) lam, sigma = get_tau_sigma(tau=lam, sigma=sigma) sigma = pt.as_tensor_variable(sigma) @@ -1937,8 +1936,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 = pt.as_tensor_variable(alpha) + m = pt.as_tensor_variable(m) return super().dist([alpha, m], **kwargs) @@ -2045,8 +2044,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 = pt.as_tensor_variable(alpha) + beta = pt.as_tensor_variable(beta) return super().dist([alpha, beta], **kwargs) @@ -2126,7 +2125,7 @@ class HalfCauchy(PositiveContinuous): @classmethod def dist(cls, beta, *args, **kwargs): - beta = pt.as_tensor_variable(floatX(beta)) + beta = pt.as_tensor_variable(beta) return super().dist([0.0, beta], **kwargs) def moment(rv, size, loc, beta): @@ -2232,8 +2231,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 = pt.as_tensor_variable(alpha) + beta = pt.as_tensor_variable(beta) # PyTensor gamma op is parametrized in terms of scale (1/beta) scale = pt.reciprocal(beta) return super().dist([alpha, scale], **kwargs) @@ -2340,8 +2339,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 = pt.as_tensor_variable(alpha) + beta = pt.as_tensor_variable(beta) return super().dist([alpha, beta], **kwargs) @@ -2523,8 +2522,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 = pt.as_tensor_variable(alpha) + beta = pt.as_tensor_variable(beta) return super().dist([alpha, beta], *args, **kwargs) @@ -2654,7 +2653,7 @@ class HalfStudentT(PositiveContinuous): @classmethod def dist(cls, nu, sigma=None, lam=None, *args, **kwargs): - nu = pt.as_tensor_variable(floatX(nu)) + nu = pt.as_tensor_variable(nu) lam, sigma = get_tau_sigma(lam, sigma) sigma = pt.as_tensor_variable(sigma) @@ -2775,9 +2774,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 = pt.as_tensor_variable(mu) + sigma = pt.as_tensor_variable(sigma) + nu = pt.as_tensor_variable(nu) return super().dist([mu, sigma, nu], *args, **kwargs) @@ -2880,8 +2879,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 = pt.as_tensor_variable(mu) + kappa = pt.as_tensor_variable(kappa) return super().dist([mu, kappa], *args, **kwargs) def moment(rv, size, mu, kappa): @@ -2984,8 +2983,8 @@ 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)) + alpha = pt.as_tensor_variable(alpha) + mu = pt.as_tensor_variable(mu) tau = pt.as_tensor_variable(tau) sigma = pt.as_tensor_variable(sigma) @@ -3073,9 +3072,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 = pt.as_tensor_variable(lower) + upper = pt.as_tensor_variable(upper) + c = pt.as_tensor_variable(c) return super().dist([lower, c, upper], *args, **kwargs) @@ -3199,8 +3198,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 = pt.as_tensor_variable(mu) + beta = pt.as_tensor_variable(beta) return super().dist([mu, beta], **kwargs) @@ -3319,8 +3318,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 = pt.as_tensor_variable(b) + sigma = pt.as_tensor_variable(sigma) return super().dist([b, sigma], *args, **kwargs) @@ -3418,8 +3417,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 = pt.as_tensor_variable(mu) + s = pt.as_tensor_variable(s) return super().dist([mu, s], *args, **kwargs) def moment(rv, size, mu, s): @@ -3523,7 +3522,7 @@ class LogitNormal(UnitContinuous): @classmethod def dist(cls, mu=0, sigma=None, tau=None, **kwargs): - mu = pt.as_tensor_variable(floatX(mu)) + mu = pt.as_tensor_variable(mu) tau, sigma = get_tau_sigma(tau=tau, sigma=sigma) sigma = pt.as_tensor_variable(sigma) tau = pt.as_tensor_variable(tau) @@ -3654,9 +3653,9 @@ def dist(cls, x_points, pdf_points, *args, **kwargs): cdf_points = interp.antiderivative()(x_points) / Z pdf_points = pdf_points / Z - x_points = pt.constant(floatX(x_points)) - pdf_points = pt.constant(floatX(pdf_points)) - cdf_points = pt.constant(floatX(cdf_points)) + x_points = pt.constant(x_points) + pdf_points = pt.constant(pdf_points) + cdf_points = pt.constant(cdf_points) # lower = pt.as_tensor_variable(x_points[0]) # upper = pt.as_tensor_variable(x_points[-1]) @@ -3693,7 +3692,7 @@ def logp(value, x_points, pdf_points, cdf_points): def interpolated_default_transform(op, rv): def transform_params(*params): _, _, _, x_points, _, _ = params - return floatX(x_points[0]), floatX(x_points[-1]) + return x_points[0], x_points[-1] return transforms.Interval(bounds_fn=transform_params) @@ -3766,8 +3765,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 = pt.as_tensor_variable(mu) + sigma = pt.as_tensor_variable(sigma) return super().dist([mu, sigma], *args, **kwargs) @@ -3862,9 +3861,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 = pt.as_tensor_variable(x) + h = pt.as_tensor_variable(h) + z = pt.as_tensor_variable(z) bshape = broadcast_shape(x, h, z) shape = [None] * len(bshape) return Apply(self, [x, h, z], [pt.TensorType(pytensor.config.floatX, shape)()]) @@ -3965,8 +3964,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 = pt.as_tensor_variable(h) + z = pt.as_tensor_variable(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 416901ae04d..238adaef7f3 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -49,7 +49,6 @@ from pymc.distributions.shape_utils import rv_size_is_none from pymc.logprob.basic import logcdf, logp from pymc.math import sigmoid -from pymc.pytensorf import floatX, intX __all__ = [ "Binomial", @@ -125,8 +124,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 = pt.as_tensor_variable(n, dtype=int) + p = pt.as_tensor_variable(p) return super().dist([n, p], **kwargs) def moment(rv, size, n, p): @@ -233,9 +232,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 = pt.as_tensor_variable(alpha) + beta = pt.as_tensor_variable(beta) + n = pt.as_tensor_variable(n, dtype=int) return super().dist([n, alpha, beta], **kwargs) def moment(rv, size, n, alpha, beta): @@ -348,7 +347,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 = pt.as_tensor_variable(p) return super().dist([p], **kwargs) def moment(rv, size, p): @@ -457,8 +456,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 = pt.as_tensor_variable(q) + beta = pt.as_tensor_variable(beta) return super().dist([q, beta], **kwargs) def moment(rv, size, q, beta): @@ -547,7 +546,7 @@ class Poisson(Discrete): @classmethod def dist(cls, mu, *args, **kwargs): - mu = pt.as_tensor_variable(floatX(mu)) + mu = pt.as_tensor_variable(mu) return super().dist([mu], *args, **kwargs) def moment(rv, size, mu): @@ -672,8 +671,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 = pt.as_tensor_variable(n) + p = pt.as_tensor_variable(p) return super().dist([n, p], *args, **kwargs) @classmethod @@ -784,7 +783,7 @@ class Geometric(Discrete): @classmethod def dist(cls, p, *args, **kwargs): - p = pt.as_tensor_variable(floatX(p)) + p = pt.as_tensor_variable(p) return super().dist([p], *args, **kwargs) def moment(rv, size, p): @@ -885,9 +884,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 = pt.as_tensor_variable(k, dtype=int) + bad = pt.as_tensor_variable(N - k, dtype=int) + n = pt.as_tensor_variable(n, dtype=int) return super().dist([good, bad, n], *args, **kwargs) def moment(rv, size, good, bad, n): @@ -1022,8 +1021,8 @@ class DiscreteUniform(Discrete): @classmethod def dist(cls, lower, upper, *args, **kwargs): - lower = intX(pt.floor(lower)) - upper = intX(pt.floor(upper)) + lower = pt.floor(lower) + upper = pt.floor(upper) return super().dist([lower, upper], **kwargs) def moment(rv, size, lower, upper): @@ -1194,7 +1193,7 @@ class _OrderedLogistic(Categorical): @classmethod def dist(cls, eta, cutpoints, *args, **kwargs): - eta = pt.as_tensor_variable(floatX(eta)) + eta = pt.as_tensor_variable(eta) cutpoints = pt.as_tensor_variable(cutpoints) pa = sigmoid(cutpoints - pt.shape_padright(eta)) @@ -1301,7 +1300,7 @@ class _OrderedProbit(Categorical): @classmethod def dist(cls, eta, cutpoints, sigma=1, *args, **kwargs): - eta = pt.as_tensor_variable(floatX(eta)) + eta = pt.as_tensor_variable(eta) cutpoints = pt.as_tensor_variable(cutpoints) probits = pt.shape_padright(eta) - cutpoints @@ -1315,7 +1314,7 @@ def dist(cls, eta, cutpoints, sigma=1, *args, **kwargs): ], axis=-1, ) - _log_p = pt.as_tensor_variable(floatX(_log_p)) + _log_p = pt.as_tensor_variable(_log_p) p = pt.exp(_log_p) return super().dist(p, *args, **kwargs) diff --git a/pymc/distributions/mixture.py b/pymc/distributions/mixture.py index 11ac3ce2437..dc270b77045 100644 --- a/pymc/distributions/mixture.py +++ b/pymc/distributions/mixture.py @@ -38,7 +38,6 @@ from pymc.logprob.abstract import _logcdf, _logcdf_helper, _logprob from pymc.logprob.basic import logp from pymc.logprob.transforms import IntervalTransform -from pymc.pytensorf import floatX from pymc.util import check_dist_not_registered from pymc.vartypes import continuous_types, discrete_types @@ -565,7 +564,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 = pt.as_tensor_variable(nonzero_p) weights = pt.stack([1 - nonzero_p, nonzero_p], axis=-1) comp_dists = [ DiracDelta.dist(0), @@ -819,7 +818,7 @@ def _hurdle_mixture(*, name, nonzero_p, nonzero_dist, dtype, **kwargs): else: raise ValueError("dtype must be 'float' or 'int'") - nonzero_p = pt.as_tensor_variable(floatX(nonzero_p)) + nonzero_p = pt.as_tensor_variable(nonzero_p) weights = pt.stack([1 - nonzero_p, nonzero_p], axis=-1) comp_dists = [ DiracDelta.dist(zero), diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 5a24bbc7cd6..8e8d510d10e 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -71,7 +71,7 @@ from pymc.distributions.transforms import Interval, ZeroSumTransform, _default_transform from pymc.logprob.abstract import _logprob from pymc.math import kron_diag, kron_dot -from pymc.pytensorf import floatX, intX +from pymc.pytensorf import intX from pymc.util import check_dist_not_registered __all__ = [ @@ -268,8 +268,8 @@ def logp(value, mu, cov): TensorVariable """ quaddist, logdet, ok = quaddist_chol(value, mu, cov) - k = floatX(value.shape[-1]) - norm = -0.5 * k * pm.floatX(np.log(2 * np.pi)) + k = value.shape[-1].astype("floatX") + norm = -0.5 * k * np.log(2 * np.pi) return check_parameters( norm - 0.5 * quaddist - logdet, ok, @@ -372,8 +372,8 @@ def dist(cls, nu, *, Sigma=None, mu=0, 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 = pt.as_tensor_variable(nu) + mu = pt.as_tensor_variable(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]) @@ -404,7 +404,7 @@ def logp(value, nu, mu, scale): TensorVariable """ quaddist, logdet, ok = quaddist_chol(value, mu, scale) - k = floatX(value.shape[-1]) + k = value.shape[-1].astype("floatX") norm = gammaln((nu + k) / 2.0) - gammaln(nu / 2.0) - 0.5 * k * pt.log(nu * np.pi) inner = -(nu + k) / 2.0 * pt.log1p(quaddist / nu) @@ -670,9 +670,6 @@ class DirichletMultinomial(Discrete): @classmethod def dist(cls, n, a, *args, **kwargs): - n = intX(n) - a = floatX(a) - return super().dist([n, a], **kwargs) def moment(rv, size, n, a): @@ -725,9 +722,9 @@ class _OrderedMultinomial(Multinomial): @classmethod def dist(cls, eta, cutpoints, n, *args, **kwargs): - eta = pt.as_tensor_variable(floatX(eta)) + eta = pt.as_tensor_variable(eta) cutpoints = pt.as_tensor_variable(cutpoints) - n = pt.as_tensor_variable(intX(n)) + n = pt.as_tensor_variable(n, dtype=int) pa = sigmoid(cutpoints - pt.shape_padright(eta)) p_cum = pt.concatenate( @@ -951,7 +948,7 @@ 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)) + V = pt.as_tensor_variable(V) warnings.warn( "The Wishart distribution can currently not be used " @@ -1181,8 +1178,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 = pt.as_tensor_variable(n, dtype=int) + eta = pt.as_tensor_variable(eta) if not ( isinstance(sd_dist, Variable) @@ -1544,8 +1541,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 = pt.as_tensor_variable(n).astype(int) + eta = pt.as_tensor_variable(eta) return super().dist([n, eta], **kwargs) def moment(rv, *args): @@ -1600,7 +1597,7 @@ def logp(value, n, eta): @_default_transform.register(_LKJCorr) def lkjcorr_default_transform(op, rv): - return MultivariateIntervalTransform(floatX(-1.0), floatX(1.0)) + return MultivariateIntervalTransform(-1.0, 1.0) class LKJCorr: @@ -1865,7 +1862,7 @@ def dist( # Broadcasting mu mu = pt.extra_ops.broadcast_to(mu, shape=dist_shape) - mu = pt.as_tensor_variable(floatX(mu)) + mu = pt.as_tensor_variable(mu) return super().dist([mu, rowchol_cov, colchol_cov], **kwargs) @@ -1909,7 +1906,7 @@ def logp(value, mu, rowchol, colchol): m = rowchol.shape[0] n = colchol.shape[0] - norm = -0.5 * m * n * pm.floatX(np.log(2 * np.pi)) + norm = -0.5 * m * n * np.log(2 * np.pi) return norm - 0.5 * trquaddist - m * half_collogdet - n * half_rowlogdet @@ -2120,13 +2117,13 @@ 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 = pt.as_tensor_variable(mu) - W = pytensor.sparse.as_sparse_or_tensor_variable(floatX(W)) + W = pytensor.sparse.as_sparse_or_tensor_variable(W) if not W.ndim == 2: raise ValueError("W must be a matrix (ndim=2).") - sparse = isinstance(W, pytensor.sparse.SparseVariable) + sparse = isinstance(W.type, pytensor.sparse.SparseTensorType) msg = "W must be a symmetric adjacency matrix." if sparse: abs_diff = pytensor.sparse.basic.mul(pytensor.sparse.sign(W - W.T), W - W.T) @@ -2134,9 +2131,9 @@ 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)) + tau = pt.as_tensor_variable(tau) - alpha = pt.as_tensor_variable(floatX(alpha)) + alpha = pt.as_tensor_variable(alpha) return super().make_node(rng, size, dtype, mu, W, alpha, tau) @@ -2447,8 +2444,8 @@ def dist(cls, W, sigma=1, zero_sum_stdev=0.001, **kwargs): N = pt.shape(W)[0] N = pt.as_tensor_variable(N) - sigma = pt.as_tensor_variable(floatX(sigma)) - zero_sum_stdev = pt.as_tensor_variable(floatX(zero_sum_stdev)) + sigma = pt.as_tensor_variable(sigma) + zero_sum_stdev = pt.as_tensor_variable(zero_sum_stdev) return super().dist([W, node1, node2, N, sigma, zero_sum_stdev], **kwargs) @@ -2557,8 +2554,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 = pt.as_tensor_variable(alpha) + K = pt.as_tensor_variable(K, dtype=int) return super().dist([alpha, K], **kwargs) @@ -2732,7 +2729,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 = pt.as_tensor_variable(sigma) if not all(sigma.type.broadcastable[-n_zerosum_axes:]): raise ValueError("sigma must have length one across the zero-sum axes") @@ -2830,8 +2827,8 @@ def zerosumnormal_logp(op, values, normal_dist, sigma, support_shape, **kwargs): *_, sigma = normal_dist.owner.inputs _deg_free_support_shape = pt.inc_subtensor(shape[-n_zerosum_axes:], -1) - _full_size = pm.floatX(pt.prod(shape)) - _degrees_of_freedom = pm.floatX(pt.prod(_deg_free_support_shape)) + _full_size = pt.prod(shape).astype("floatX") + _degrees_of_freedom = pt.prod(_deg_free_support_shape).astype("floatX") zerosums = [ pt.all(pt.isclose(pt.mean(value, axis=-axis - 1), 0, atol=1e-9)) diff --git a/pymc/distributions/simulator.py b/pymc/distributions/simulator.py index cafc65cf4bf..a2e6357f18e 100644 --- a/pymc/distributions/simulator.py +++ b/pymc/distributions/simulator.py @@ -25,7 +25,6 @@ from pymc.distributions.distribution import Distribution, _moment from pymc.logprob.abstract import _logprob -from pymc.pytensorf import floatX __all__ = ["Simulator"] @@ -192,7 +191,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 = pt.as_tensor_variable(epsilon) if params is None: params = unnamed_params diff --git a/pymc/distributions/timeseries.py b/pymc/distributions/timeseries.py index 6bf653c3561..e3e0de28d41 100644 --- a/pymc/distributions/timeseries.py +++ b/pymc/distributions/timeseries.py @@ -43,7 +43,7 @@ from pymc.exceptions import NotConstantValueError from pymc.logprob.abstract import _logprob from pymc.logprob.basic import logp -from pymc.pytensorf import constant_fold, floatX, intX +from pymc.pytensorf import constant_fold, intX from pymc.util import check_dist_not_registered __all__ = [ @@ -495,7 +495,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(pt.as_tensor_variable(rho)) ar_order = cls._get_ar_order(rhos=rhos, constant=constant, ar_order=ar_order) steps = get_support_shape_1d( support_shape=steps, @@ -522,8 +522,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 = pt.as_tensor_variable(sigma) + rhos = pt.atleast_1d(pt.as_tensor_variable(rho)) if "init" in kwargs: warnings.warn( @@ -916,7 +916,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 = pt.as_tensor_variable(dt) steps = get_support_shape_1d( support_shape=steps, shape=None, # Shape will be checked in `cls.dist` @@ -935,7 +935,7 @@ 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)) + dt = pt.as_tensor_variable(dt) sde_pars = [pt.as_tensor_variable(x) for x in sde_pars] if init_dist is not None: diff --git a/pymc/distributions/transforms.py b/pymc/distributions/transforms.py index 49c053e95f3..aeedceedd3f 100644 --- a/pymc/distributions/transforms.py +++ b/pymc/distributions/transforms.py @@ -25,8 +25,6 @@ from pytensor.graph import Op from pytensor.tensor import TensorVariable -import pymc as pm - from pymc.logprob.transforms import ( ChainedTransform, CircularTransform, @@ -284,7 +282,7 @@ def __init__(self, zerosum_axes): @staticmethod def extend_axis(array, axis): - n = pm.floatX(array.shape[axis] + 1) + n = (array.shape[axis] + 1).astype("floatX") sum_vals = array.sum(axis, keepdims=True) norm = sum_vals / (pt.sqrt(n) + n) fill_val = norm - sum_vals / pt.sqrt(n) @@ -296,7 +294,7 @@ def extend_axis(array, axis): def extend_axis_rev(array, axis): normalized_axis = normalize_axis_tuple(axis, array.ndim)[0] - n = pm.floatX(array.shape[normalized_axis]) + n = array.shape[normalized_axis].astype("floatX") last = pt.take(array, [-1], axis=normalized_axis) sum_vals = -last * pt.sqrt(n) diff --git a/tests/distributions/test_mixture.py b/tests/distributions/test_mixture.py index 6d01efcec9a..767fa618b6e 100644 --- a/tests/distributions/test_mixture.py +++ b/tests/distributions/test_mixture.py @@ -19,7 +19,7 @@ import pytest import scipy.stats as st -from numpy.testing import assert_allclose +from numpy.testing import assert_allclose, assert_almost_equal from pytensor import tensor as pt from pytensor.tensor import TensorVariable from pytensor.tensor.random.op import RandomVariable @@ -1638,7 +1638,7 @@ def test_hurdle_lognormal_graph(self): ], ) def test_hurdle_logp_at_zero(self, dist, psi, non_psi_args): - assert logp(dist(psi=psi, **non_psi_args), 0).eval() == np.log(1 - psi) + assert_almost_equal(logp(dist(psi=psi, **non_psi_args), 0).eval(), np.log(1 - psi)) @pytest.mark.parametrize( "dist, psi, non_psi_args", diff --git a/tests/distributions/test_truncated.py b/tests/distributions/test_truncated.py index f9b62a5c832..cbe50b13f67 100644 --- a/tests/distributions/test_truncated.py +++ b/tests/distributions/test_truncated.py @@ -412,7 +412,10 @@ def test_truncated_gamma(): logp_scipy[x > upper] = -np.inf gamma_trunc_pymc = Truncated.dist( - Gamma.dist(alpha=alpha, beta=beta), + Gamma.dist( + alpha=pt.as_tensor_variable(alpha).astype("floatX"), + beta=pt.as_tensor_variable(beta).astype("floatX"), + ), upper=upper, ) logp_pymc = logp(gamma_trunc_pymc, x).eval() diff --git a/tests/test_initial_point.py b/tests/test_initial_point.py index ccf10b3fd96..8e8ac3018ca 100644 --- a/tests/test_initial_point.py +++ b/tests/test_initial_point.py @@ -13,6 +13,7 @@ # limitations under the License. import cloudpickle import numpy as np +import numpy.testing as npt import pytensor import pytensor.tensor as pt import pytest @@ -48,7 +49,9 @@ def test_new_warnings(self): with pytest.warns(FutureWarning, match="`testval` argument is deprecated"): rv = pm.Uniform("u", 0, 1, testval=0.75) initial_point = pmodel.initial_point(random_seed=0) - assert initial_point["u_interval__"] == transform_fwd(rv, 0.75, model=pmodel) + npt.assert_allclose( + initial_point["u_interval__"], transform_fwd(rv, 0.75, model=pmodel) + ) assert not hasattr(rv.tag, "test_value") pass diff --git a/tests/variational/test_inference.py b/tests/variational/test_inference.py index 6f04819ba1d..2e0a3c18871 100644 --- a/tests/variational/test_inference.py +++ b/tests/variational/test_inference.py @@ -81,12 +81,15 @@ def simple_model_data(use_minibatch): def simple_model(simple_model_data): with pm.Model() as model: mu_ = pm.Normal( - "mu", mu=simple_model_data["mu0"], sigma=simple_model_data["sigma0"], initval=0 + "mu", + mu=pt.as_tensor_variable(simple_model_data["mu0"]).astype("floatX"), + sigma=pt.as_tensor_variable(simple_model_data["sigma0"]).astype("floatX"), + initval=0, ) pm.Normal( "x", mu=mu_, - sigma=simple_model_data["sigma"], + sigma=pt.as_tensor_variable(simple_model_data["sigma"]).astype("floatX"), observed=simple_model_data["data"], total_size=simple_model_data["n"], )