diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index 179bae25f..7a33574ab 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -13,7 +13,10 @@ # limitations under the License. import warnings +from typing import TypeAlias + import numpy as np +import numpy.typing as npt import pytensor.tensor as pt from pytensor.tensor import TensorConstant @@ -30,6 +33,7 @@ uniform, ) from pytensor.tensor.random.utils import normalize_size_param +from pytensor.tensor.variable import TensorVariable from scipy import stats import pymc as pm @@ -46,10 +50,11 @@ normal_lccdf, normal_lcdf, ) -from pymc.distributions.distribution import Discrete, SymbolicRandomVariable +from pymc.distributions.distribution import DIST_PARAMETER_TYPES, Discrete, SymbolicRandomVariable from pymc.distributions.shape_utils import implicit_size_from_params, rv_size_is_none from pymc.logprob.basic import logcdf, logp from pymc.math import sigmoid +from pymc.pytensorf import normalize_rng_param __all__ = [ "Binomial", @@ -66,7 +71,7 @@ "OrderedProbit", ] -from pymc.pytensorf import normalize_rng_param +DISCRETE_DIST_PARAMETER_TYPES: TypeAlias = npt.NDArray[np.int_] | int | TensorVariable class Binomial(Discrete): @@ -118,7 +123,14 @@ class Binomial(Discrete): rv_op = binomial @classmethod - def dist(cls, n, p=None, logit_p=None, *args, **kwargs): + def dist( + cls, + n: DISCRETE_DIST_PARAMETER_TYPES, + p: DIST_PARAMETER_TYPES | None = None, + logit_p: DIST_PARAMETER_TYPES | None = None, + *args, + **kwargs, + ): if p is not None and logit_p is not None: raise ValueError("Incompatible parametrization. Can't specify both p and logit_p.") elif p is None and logit_p is None: @@ -234,7 +246,14 @@ def BetaBinom(a, b, n, x): rv_op = betabinom @classmethod - def dist(cls, alpha, beta, n, *args, **kwargs): + def dist( + cls, + alpha: DIST_PARAMETER_TYPES, + beta: DIST_PARAMETER_TYPES, + n: DISCRETE_DIST_PARAMETER_TYPES, + *args, + **kwargs, + ): alpha = pt.as_tensor_variable(alpha) beta = pt.as_tensor_variable(beta) n = pt.as_tensor_variable(n, dtype=int) @@ -341,7 +360,13 @@ class Bernoulli(Discrete): rv_op = bernoulli @classmethod - def dist(cls, p=None, logit_p=None, *args, **kwargs): + def dist( + cls, + p: DIST_PARAMETER_TYPES | None = None, + logit_p: DIST_PARAMETER_TYPES | None = None, + *args, + **kwargs, + ): if p is not None and logit_p is not None: raise ValueError("Incompatible parametrization. Can't specify both p and logit_p.") elif p is None and logit_p is None: @@ -465,7 +490,9 @@ def DiscreteWeibull(q, b, x): rv_op = DiscreteWeibullRV.rv_op @classmethod - def dist(cls, q, beta, *args, **kwargs): + def dist(cls, q: DIST_PARAMETER_TYPES, beta: DIST_PARAMETER_TYPES, *args, **kwargs): + q = pt.as_tensor_variable(q) + beta = pt.as_tensor_variable(beta) return super().dist([q, beta], **kwargs) def support_point(rv, size, q, beta): @@ -553,7 +580,7 @@ class Poisson(Discrete): rv_op = poisson @classmethod - def dist(cls, mu, *args, **kwargs): + def dist(cls, mu: DIST_PARAMETER_TYPES, *args, **kwargs): mu = pt.as_tensor_variable(mu) return super().dist([mu], *args, **kwargs) @@ -677,7 +704,15 @@ def NegBinom(a, m, x): rv_op = nbinom @classmethod - def dist(cls, mu=None, alpha=None, p=None, n=None, *args, **kwargs): + def dist( + cls, + mu: DIST_PARAMETER_TYPES | None = None, + alpha: DIST_PARAMETER_TYPES | None = None, + p: DIST_PARAMETER_TYPES | None = None, + n: DIST_PARAMETER_TYPES | None = None, + *args, + **kwargs, + ): n, p = cls.get_n_p(mu=mu, alpha=alpha, p=p, n=n) n = pt.as_tensor_variable(n) p = pt.as_tensor_variable(p) @@ -790,7 +825,7 @@ class Geometric(Discrete): rv_op = geometric @classmethod - def dist(cls, p, *args, **kwargs): + def dist(cls, p: DIST_PARAMETER_TYPES, *args, **kwargs): p = pt.as_tensor_variable(p) return super().dist([p], *args, **kwargs) @@ -891,7 +926,14 @@ class HyperGeometric(Discrete): rv_op = hypergeometric @classmethod - def dist(cls, N, k, n, *args, **kwargs): + def dist( + cls, + N: DISCRETE_DIST_PARAMETER_TYPES, + k: DISCRETE_DIST_PARAMETER_TYPES, + n: DISCRETE_DIST_PARAMETER_TYPES, + *args, + **kwargs, + ): 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) @@ -1027,7 +1069,13 @@ class DiscreteUniform(Discrete): rv_op = discrete_uniform @classmethod - def dist(cls, lower, upper, *args, **kwargs): + def dist( + cls, + lower: DISCRETE_DIST_PARAMETER_TYPES, + upper: DISCRETE_DIST_PARAMETER_TYPES, + *args, + **kwargs, + ): lower = pt.floor(lower) upper = pt.floor(upper) return super().dist([lower, upper], **kwargs) @@ -1123,7 +1171,12 @@ class Categorical(Discrete): rv_op = categorical @classmethod - def dist(cls, p=None, logit_p=None, **kwargs): + def dist( + cls, + p: np.ndarray | None = None, + logit_p: float | None = None, + **kwargs, + ): if p is not None and logit_p is not None: raise ValueError("Incompatible parametrization. Can't specify both p and logit_p.") elif p is None and logit_p is None: @@ -1261,7 +1314,7 @@ def __new__(cls, name, eta, cutpoints, compute_p=True, **kwargs): return out_rv @classmethod - def dist(cls, eta, cutpoints, **kwargs): + def dist(cls, eta: DIST_PARAMETER_TYPES, cutpoints: DIST_PARAMETER_TYPES, **kwargs): p = cls.compute_p(eta, cutpoints) return Categorical.dist(p=p, **kwargs)