Add type hints to dist() in discrete distributions

Author: johanbog

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)