Fix some type hinting to help with migrating Distribution (pymc-devs#7484)

Author: thomasaarholt

pymc/distributions/continuous.py

@@ -51,10 +51,10 @@
 )
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.random.utils import normalize_size_param
-from pytensor.tensor.variable import TensorConstant
+from pytensor.tensor.variable import TensorConstant, TensorVariable
 
 from pymc.logprob.abstract import _logprob_helper
-from pymc.logprob.basic import icdf
+from pymc.logprob.basic import TensorLike, icdf
 from pymc.pytensorf import normalize_rng_param
 
 try:
@@ -148,7 +148,7 @@ class BoundedContinuous(Continuous):
     """Base class for bounded continuous distributions."""
 
     # Indices of the arguments that define the lower and upper bounds of the distribution
-    bound_args_indices: list[int] | None = None
+    bound_args_indices: tuple[int | None, int | None] | None = None
 
 
 @_default_transform.register(PositiveContinuous)
@@ -210,7 +210,9 @@ def assert_negative_support(var, label, distname, value=-1e-6):
     return Assert(msg)(var, pt.all(pt.ge(var, 0.0)))
 
 
-def get_tau_sigma(tau=None, sigma=None):
+def get_tau_sigma(
+    tau: TensorLike | None = None, sigma: TensorLike | None = None
+) -> tuple[TensorVariable, TensorVariable]:
     r"""
     Find precision and standard deviation.
 
@@ -239,13 +241,14 @@ def get_tau_sigma(tau=None, sigma=None):
         sigma = pt.as_tensor_variable(1.0)
         tau = pt.as_tensor_variable(1.0)
     elif tau is None:
+        assert sigma is not None  # Just for type checker
         sigma = pt.as_tensor_variable(sigma)
         # Keep tau negative, if sigma was negative, so that it will
         # fail when used
         tau = (sigma**-2.0) * pt.sign(sigma)
     else:
         tau = pt.as_tensor_variable(tau)
-        # Keep tau negative, if sigma was negative, so that it will
+        # Keep sigma negative, if tau was negative, so that it will
         # fail when used
         sigma = pt.abs(tau) ** -0.5 * pt.sign(tau)
 

pymc/distributions/distribution.py

@@ -21,7 +21,7 @@
 from abc import ABCMeta
 from collections.abc import Callable, Sequence
 from functools import singledispatch
-from typing import TypeAlias
+from typing import Any, TypeAlias
 
 import numpy as np
 
@@ -423,8 +423,12 @@ def change_symbolic_rv_size(op: SymbolicRandomVariable, rv, new_size, expand) ->
 class Distribution(metaclass=DistributionMeta):
     """Statistical distribution."""
 
-    rv_op: [RandomVariable, SymbolicRandomVariable] = None
-    rv_type: MetaType = None
+    # rv_op and _type are set to None via the DistributionMeta.__new__
+    # if not specified as class attributes in subclasses of Distribution.
+    # rv_op can either be a class (see the Normal class) or a method
+    # (see the Censored class), both callable to return a TensorVariable.
+    rv_op: Any = None
+    rv_type: MetaType | None = None
 
     def __new__(
         cls,