From a11d9039744cd0f0a8717b7af088c52687ce7817 Mon Sep 17 00:00:00 2001 From: Johannes Bogen Date: Sat, 30 Sep 2023 13:47:28 -0700 Subject: [PATCH 1/3] Add type hints to dist() in discrete distributions --- pymc/distributions/discrete.py | 85 ++++++++++++++++++++++++++++------ 1 file changed, 72 insertions(+), 13 deletions(-) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index c18a88690..f9791b606 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -13,7 +13,10 @@ # limitations under the License. import warnings +from typing import Optional, TypeAlias, Union + import numpy as np +import numpy.typing as npt import pytensor.tensor as pt from pytensor.tensor import TensorConstant @@ -29,6 +32,7 @@ nbinom, poisson, ) +from pytensor.tensor.variable import TensorVariable from scipy import stats import pymc as pm @@ -45,7 +49,7 @@ normal_lccdf, normal_lcdf, ) -from pymc.distributions.distribution import Discrete +from pymc.distributions.distribution import DIST_PARAMETER_TYPES, Discrete from pymc.distributions.shape_utils import rv_size_is_none from pymc.logprob.basic import logcdf, logp from pymc.math import sigmoid @@ -65,6 +69,8 @@ "OrderedProbit", ] +DISCRETE_DIST_PARAMETER_TYPES: TypeAlias = Union[npt.NDArray[np.int_], int, TensorVariable] + class Binomial(Discrete): R""" @@ -115,7 +121,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: Optional[DIST_PARAMETER_TYPES] = None, + logit_p: Optional[DIST_PARAMETER_TYPES] = 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: @@ -231,7 +244,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) @@ -338,7 +358,13 @@ class Bernoulli(Discrete): rv_op = bernoulli @classmethod - def dist(cls, p=None, logit_p=None, *args, **kwargs): + def dist( + cls, + p: Optional[DIST_PARAMETER_TYPES] = None, + logit_p: Optional[DIST_PARAMETER_TYPES] = 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: @@ -455,7 +481,7 @@ def DiscreteWeibull(q, b, x): rv_op = discrete_weibull @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) @@ -545,7 +571,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) @@ -669,7 +695,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: Optional[DIST_PARAMETER_TYPES] = None, + alpha: Optional[DIST_PARAMETER_TYPES] = None, + p: Optional[DIST_PARAMETER_TYPES] = None, + n: Optional[DIST_PARAMETER_TYPES] = 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) @@ -782,7 +816,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) @@ -883,7 +917,14 @@ class HyperGeometric(Discrete): rv_op = hypergeometric @classmethod - def dist(cls, N, k, n, *args, **kwargs): + def dist( + cls, + N: Optional[DISCRETE_DIST_PARAMETER_TYPES], + k: Optional[DISCRETE_DIST_PARAMETER_TYPES], + n: Optional[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) @@ -1020,7 +1061,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) @@ -1116,7 +1163,12 @@ class Categorical(Discrete): rv_op = categorical @classmethod - def dist(cls, p=None, logit_p=None, **kwargs): + def dist( + cls, + p: Optional[np.ndarray] = None, + logit_p: Optional[float] = 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: @@ -1192,7 +1244,7 @@ class _OrderedLogistic(Categorical): rv_op = categorical @classmethod - def dist(cls, eta, cutpoints, *args, **kwargs): + def dist(cls, eta: DIST_PARAMETER_TYPES, cutpoints: DIST_PARAMETER_TYPES, *args, **kwargs): eta = pt.as_tensor_variable(eta) cutpoints = pt.as_tensor_variable(cutpoints) @@ -1299,7 +1351,14 @@ class _OrderedProbit(Categorical): rv_op = categorical @classmethod - def dist(cls, eta, cutpoints, sigma=1, *args, **kwargs): + def dist( + cls, + eta: DIST_PARAMETER_TYPES, + cutpoints: DIST_PARAMETER_TYPES, + sigma: DIST_PARAMETER_TYPES = 1.0, + *args, + **kwargs, + ): eta = pt.as_tensor_variable(eta) cutpoints = pt.as_tensor_variable(cutpoints) From 50123a144570bdf0bc310bcd1b20f1e4c117a98e Mon Sep 17 00:00:00 2001 From: Johannes Bogen Date: Sat, 30 Sep 2023 19:03:50 -0700 Subject: [PATCH 2/3] fix wrong import --- pymc/distributions/discrete.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index f9791b606..e86db6393 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -13,7 +13,7 @@ # limitations under the License. import warnings -from typing import Optional, TypeAlias, Union +from typing import Optional, Union import numpy as np import numpy.typing as npt @@ -34,6 +34,7 @@ ) from pytensor.tensor.variable import TensorVariable from scipy import stats +from typing_extensions import TypeAlias import pymc as pm From 6f6b416d52357ea05cb29859f5d28472ef59b33b Mon Sep 17 00:00:00 2001 From: Johannes Bogen Date: Fri, 1 Mar 2024 18:17:23 -0800 Subject: [PATCH 3/3] fix non-optional argument set as optional --- pymc/distributions/discrete.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index e86db6393..23003691b 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -920,9 +920,9 @@ class HyperGeometric(Discrete): @classmethod def dist( cls, - N: Optional[DISCRETE_DIST_PARAMETER_TYPES], - k: Optional[DISCRETE_DIST_PARAMETER_TYPES], - n: Optional[DISCRETE_DIST_PARAMETER_TYPES], + N: DISCRETE_DIST_PARAMETER_TYPES, + k: DISCRETE_DIST_PARAMETER_TYPES, + n: DISCRETE_DIST_PARAMETER_TYPES, *args, **kwargs, ):