Add more type hints to distribution parameter

Author: iykat

pymc/distributions/continuous.py

@@ -959,7 +959,7 @@ def dist(
         mu: Optional[DIST_PARAMETER_TYPES] = None,
         lam: Optional[DIST_PARAMETER_TYPES] = None,
         phi: Optional[DIST_PARAMETER_TYPES] = None,
-        alpha=0.0,
+        alpha: Optional[DIST_PARAMETER_TYPES] = 0.0,
         **kwargs,
     ):
         mu, lam, phi = cls.get_mu_lam_phi(mu, lam, phi)
@@ -1128,7 +1128,16 @@ class Beta(UnitContinuous):
     rv_op = pytensor.tensor.random.beta
 
     @classmethod
-    def dist(cls, alpha=None, beta=None, mu=None, sigma=None, nu=None, *args, **kwargs):
+    def dist(
+        cls,
+        alpha: Optional[DIST_PARAMETER_TYPES] = None,
+        beta: Optional[DIST_PARAMETER_TYPES] = None,
+        mu: Optional[DIST_PARAMETER_TYPES] = None,
+        sigma: Optional[DIST_PARAMETER_TYPES] = None,
+        nu: Optional[DIST_PARAMETER_TYPES] = 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))
@@ -1256,7 +1265,7 @@ class Kumaraswamy(UnitContinuous):
     rv_op = kumaraswamy
 
     @classmethod
-    def dist(cls, a, b, *args, **kwargs):
+    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))
 
@@ -1342,7 +1351,7 @@ class Exponential(PositiveContinuous):
     rv_op = exponential
 
     @classmethod
-    def dist(cls, lam, *args, **kwargs):
+    def dist(cls, lam: DIST_PARAMETER_TYPES, *args, **kwargs):
         lam = pt.as_tensor_variable(floatX(lam))
 
         # PyTensor exponential op is parametrized in terms of mu (1/lam)