Update dist parameter hints (#5315)

Author: canyon289

pymc/distributions/continuous.py

@@ -85,7 +85,7 @@ def polyagamma_cdf(*args, **kwargs):
     normal_lcdf,
     zvalue,
 )
-from pymc.distributions.distribution import Continuous
+from pymc.distributions.distribution import DIST_PARAMETER_TYPES, Continuous
 from pymc.distributions.shape_utils import rv_size_is_none
 from pymc.math import invlogit, logdiffexp, logit
 from pymc.util import UNSET
@@ -692,12 +692,12 @@ class TruncatedNormal(BoundedContinuous):
     @classmethod
     def dist(
         cls,
-        mu: Optional[Union[float, np.ndarray]] = None,
-        sigma: Optional[Union[float, np.ndarray]] = None,
-        tau: Optional[Union[float, np.ndarray]] = None,
-        sd: Optional[Union[float, np.ndarray]] = None,
-        lower: Optional[Union[float, np.ndarray]] = None,
-        upper: Optional[Union[float, np.ndarray]] = None,
+        mu: Optional[DIST_PARAMETER_TYPES] = None,
+        sigma: Optional[DIST_PARAMETER_TYPES] = None,
+        tau: Optional[DIST_PARAMETER_TYPES] = None,
+        sd: Optional[DIST_PARAMETER_TYPES] = None,
+        lower: Optional[DIST_PARAMETER_TYPES] = None,
+        upper: Optional[DIST_PARAMETER_TYPES] = None,
         transform: str = "auto",
         *args,
         **kwargs,

pymc/distributions/distribution.py

@@ -61,6 +61,8 @@
     "NoDistribution",
 ]
 
+DIST_PARAMETER_TYPES = Union[np.ndarray, int, float, TensorVariable]
+
 vectorized_ppc = contextvars.ContextVar(
     "vectorized_ppc", default=None
 )  # type: contextvars.ContextVar[Optional[Callable]]