Add return type hints (#5819)

* Add return type hints
* Add type hints for rv arguments
* Fix more type issues

Closes #4880

Co-authored-by: Ricardo Vieira <[email protected]>

Author: michaelosthege

pymc/aesaraf.py

@@ -52,6 +52,10 @@
 from aesara.scalar.basic import Cast
 from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.op import RandomVariable
+from aesara.tensor.random.var import (
+    RandomGeneratorSharedVariable,
+    RandomStateSharedVariable,
+)
 from aesara.tensor.shape import SpecifyShape
 from aesara.tensor.sharedvar import SharedVariable
 from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
@@ -60,9 +64,7 @@
 from pymc.exceptions import ShapeError
 from pymc.vartypes import continuous_types, isgenerator, typefilter
 
-PotentialShapeType = Union[
-    int, np.ndarray, Tuple[Union[int, Variable], ...], List[Union[int, Variable]], Variable
-]
+PotentialShapeType = Union[int, np.ndarray, Sequence[Union[int, Variable]], TensorVariable]
 
 
 __all__ = [
@@ -165,6 +167,7 @@ def change_rv_size(
         new_size = (new_size,)
 
     # Extract the RV node that is to be resized, together with its inputs, name and tag
+    assert rv.owner.op is not None
     if isinstance(rv.owner.op, SpecifyShape):
         rv = rv.owner.inputs[0]
     rv_node = rv.owner
@@ -894,18 +897,14 @@ def local_check_parameter_to_ninf_switch(fgraph, node):
 )
 
 
-def find_rng_nodes(variables: Iterable[TensorVariable]):
+def find_rng_nodes(
+    variables: Iterable[Variable],
+) -> List[Union[RandomStateSharedVariable, RandomGeneratorSharedVariable]]:
     """Return RNG variables in a graph"""
     return [
         node
         for node in graph_inputs(variables)
-        if isinstance(
-            node,
-            (
-                at.random.var.RandomStateSharedVariable,
-                at.random.var.RandomGeneratorSharedVariable,
-            ),
-        )
+        if isinstance(node, (RandomStateSharedVariable, RandomGeneratorSharedVariable))
     ]
 
 
@@ -921,6 +920,7 @@ def reseed_rngs(
         np.random.PCG64(sub_seed) for sub_seed in np.random.SeedSequence(seed).spawn(len(rngs))
     ]
     for rng, bit_generator in zip(rngs, bit_generators):
+        new_rng: Union[np.random.RandomState, np.random.Generator]
         if isinstance(rng, at.random.var.RandomStateSharedVariable):
             new_rng = np.random.RandomState(bit_generator)
         else:
@@ -980,6 +980,9 @@ def compile_pymc(
         and isinstance(var.owner.op, (RandomVariable, MeasurableVariable))
         and var not in inputs
     ):
+        # All nodes in `vars_between(inputs, outputs)` have owners.
+        # But mypy doesn't know, so we just assert it:
+        assert random_var.owner.op is not None
         if isinstance(random_var.owner.op, RandomVariable):
             rng = random_var.owner.inputs[0]
             if not hasattr(rng, "default_update"):

pymc/distributions/distribution.py

@@ -198,8 +198,8 @@ def __new__(
         total_size=None,
         transform=UNSET,
         **kwargs,
-    ) -> RandomVariable:
-        """Adds a RandomVariable corresponding to a PyMC distribution to the current model.
+    ) -> TensorVariable:
+        """Adds a tensor variable corresponding to a PyMC distribution to the current model.
 
         Note that all remaining kwargs must be compatible with ``.dist()``
 
@@ -231,8 +231,8 @@ def __new__(
 
         Returns
         -------
-        rv : RandomVariable
-            The created RV, registered in the Model.
+        rv : TensorVariable
+            The created random variable tensor, registered in the Model.
         """
 
         try:
@@ -296,8 +296,8 @@ def dist(
         *,
         shape: Optional[Shape] = None,
         **kwargs,
-    ) -> RandomVariable:
-        """Creates a RandomVariable corresponding to the `cls` distribution.
+    ) -> TensorVariable:
+        """Creates a tensor variable corresponding to the `cls` distribution.
 
         Parameters
         ----------
@@ -314,8 +314,8 @@ def dist(
 
         Returns
         -------
-        rv : RandomVariable
-            The created RV.
+        rv : TensorVariable
+            The created random variable tensor.
         """
         if "testval" in kwargs:
             kwargs.pop("testval")
@@ -653,8 +653,8 @@ def __new__(
         name : str
         dist_params : Tuple
             A sequence of the distribution's parameter. These will be converted into
-            Aesara tensors internally. These parameters could be other ``RandomVariable``
-            instances.
+            Aesara tensors internally. These parameters could be other ``TensorVariable``
+            instances created from , optionally created via ``RandomVariable`` ``Op``s.
         logp : Optional[Callable]
             A callable that calculates the log density of some given observed ``value``
             conditioned on certain distribution parameter values. It must have the

pymc/distributions/logprob.py

@@ -39,7 +39,9 @@
 from pymc.aesaraf import floatX
 
 
-def _get_scaling(total_size: Optional[Union[int, Sequence[int]]], shape, ndim: int):
+def _get_scaling(
+    total_size: Optional[Union[int, Sequence[int]]], shape, ndim: int
+) -> TensorVariable:
     """
     Gets scaling constant for logp.
 
@@ -288,14 +290,14 @@ def logp(rv: TensorVariable, value) -> TensorVariable:
             raise NotImplementedError("PyMC could not infer logp of input variable.") from exc
 
 
-def logcdf(rv, value):
+def logcdf(rv: TensorVariable, value) -> TensorVariable:
     """Return the log-cdf graph of a Random Variable"""
 
     value = at.as_tensor_variable(value, dtype=rv.dtype)
     return logcdf_aeppl(rv, value)
 
 
-def ignore_logprob(rv):
+def ignore_logprob(rv: TensorVariable) -> TensorVariable:
     """Return a duplicated variable that is ignored when creating Aeppl logprob graphs
 
     This is used in SymbolicDistributions that use other RVs as inputs but account