More general type-hints

Author: ricardoV94

pymc/distributions/shape_utils.py

@@ -230,17 +230,17 @@ def rv_size_is_none(size: TensorVariable | Constant | None) -> bool:
 
 
 @singledispatch
-def _change_dist_size(op: Op, dist: TensorVariable, new_size, expand):
+def _change_dist_size(op: Op, dist: Variable, new_size, expand):
     raise NotImplementedError(
         f"Variable {dist} of type {op} has no _change_dist_size implementation."
     )
 
 
 def change_dist_size(
-    dist: TensorVariable,
+    dist: Variable,
     new_size: PotentialShapeType,
     expand: bool = False,
-) -> TensorVariable:
+) -> Variable:
     """Change or expand the size of a Distribution.
 
     Parameters

pymc/sampling/forward.py

@@ -42,7 +42,7 @@
 from pytensor.graph.traversal import ancestors, general_toposort, walk
 from pytensor.tensor.random.var import RandomGeneratorSharedVariable
 from pytensor.tensor.sharedvar import SharedVariable, TensorSharedVariable
-from pytensor.tensor.variable import TensorConstant, TensorVariable
+from pytensor.tensor.variable import TensorConstant
 from rich.console import Console
 from rich.progress import BarColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn
 from rich.theme import Theme
@@ -1032,10 +1032,8 @@ def vectorize_over_posterior(
         If random variables are found in the graph and `allow_rvs_in_graph` is False
     """
     # Identify which free RVs are needed to compute `outputs`
-    needed_rvs: list[TensorVariable] = [
-        cast(TensorVariable, rv)
-        for rv in ancestors(outputs, blockers=input_rvs)
-        if rv in set(input_rvs)
+    needed_rvs: list[Variable] = [
+        rv for rv in ancestors(outputs, blockers=input_rvs) if rv in set(input_rvs)
     ]
 
     # Replace needed_rvs with actual posterior samples
@@ -1044,7 +1042,7 @@ def vectorize_over_posterior(
     for rv in needed_rvs:
         posterior_samples = posterior[rv.name].data
 
-        replace_dict[rv] = pt.constant(posterior_samples.astype(rv.dtype), name=rv.name)
+        replace_dict[rv] = pt.constant(posterior_samples.astype(rv.dtype), name=rv.name)  # type: ignore[attr-defined]
 
     # Replace the rvs that remain in the graph with resized versions
     all_rvs = rvs_in_graph(outputs)