Even more type fixes

Author: michaelosthege

pymc/backends/__init__.py

@@ -81,7 +81,7 @@
 
     from pymc.backends.mcbackend import init_chain_adapters
 
-    TraceOrBackend = BaseTrace | Backend
+    TraceOrBackend: TypeAlias = BaseTrace | Backend
     RunType: TypeAlias = Run
     HAS_MCB = True
 except ImportError:

pymc/distributions/distribution.py

@@ -28,7 +28,7 @@
 from pytensor import tensor as pt
 from pytensor.compile.builders import OpFromGraph
 from pytensor.graph import FunctionGraph, clone_replace, node_rewriter
-from pytensor.graph.basic import Node, Variable, io_toposort
+from pytensor.graph.basic import Apply, Variable, io_toposort
 from pytensor.graph.features import ReplaceValidate
 from pytensor.graph.rewriting.basic import GraphRewriter, in2out
 from pytensor.graph.utils import MetaType
@@ -421,7 +421,7 @@ def __init__(
         kwargs.setdefault("strict", True)
         super().__init__(*args, **kwargs)
 
-    def update(self, node: Node) -> dict[Variable, Variable]:
+    def update(self, node: Apply) -> dict[Variable, Variable]:
         """Symbolic update expression for input random state variables
 
         Returns a dictionary with the symbolic expressions required for correct updating
@@ -430,7 +430,7 @@ def update(self, node: Node) -> dict[Variable, Variable]:
         """
         return collect_default_updates_inner_fgraph(node)
 
-    def batch_ndim(self, node: Node) -> int:
+    def batch_ndim(self, node: Apply) -> int:
         """Number of dimensions of the distribution's batch shape."""
         out_ndim = max(getattr(out.type, "ndim", 0) for out in node.outputs)
         return out_ndim - self.ndim_supp

pymc/distributions/mixture.py

@@ -18,7 +18,7 @@
 import pytensor
 import pytensor.tensor as pt
 
-from pytensor.graph.basic import Node, equal_computations
+from pytensor.graph.basic import Apply, equal_computations
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.random.utils import normalize_size_param
@@ -156,7 +156,7 @@ def _resize_components(cls, size, *components):
 
         return [change_dist_size(component, size) for component in components]
 
-    def update(self, node: Node):
+    def update(self, node: Apply):
         # Update for the internal mix_indexes RV
         return {node.inputs[0]: node.outputs[0]}
 

pymc/logprob/order.py

@@ -35,9 +35,11 @@
 #   SOFTWARE.
 
 
+from typing import cast
+
 import pytensor.tensor as pt
 
-from pytensor.graph.basic import Node
+from pytensor.graph.basic import Apply
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import node_rewriter
 from pytensor.tensor.elemwise import Elemwise
@@ -72,15 +74,15 @@ class MeasurableMaxDiscrete(Max):
 
 
 @node_rewriter([Max])
-def find_measurable_max(fgraph: FunctionGraph, node: Node) -> list[TensorVariable] | None:
+def find_measurable_max(fgraph: FunctionGraph, node: Apply) -> list[TensorVariable] | None:
     rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)
     if rv_map_feature is None:
         return None  # pragma: no cover
 
     if isinstance(node.op, MeasurableMax):
         return None  # pragma: no cover
 
-    base_var = node.inputs[0]
+    base_var = cast(TensorVariable, node.inputs[0])
 
     if base_var.owner is None:
         return None
@@ -104,6 +106,7 @@ def find_measurable_max(fgraph: FunctionGraph, node: Node) -> list[TensorVariabl
         return None
 
     # distinguish measurable discrete and continuous (because logprob is different)
+    measurable_max: Max
     if base_var.owner.op.dtype.startswith("int"):
         measurable_max = MeasurableMaxDiscrete(list(axis))
     else:
@@ -173,7 +176,7 @@ class MeasurableDiscreteMaxNeg(Max):
 
 
 @node_rewriter(tracks=[Max])
-def find_measurable_max_neg(fgraph: FunctionGraph, node: Node) -> list[TensorVariable] | None:
+def find_measurable_max_neg(fgraph: FunctionGraph, node: Apply) -> list[TensorVariable] | None:
     rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)
 
     if rv_map_feature is None:
@@ -182,7 +185,7 @@ def find_measurable_max_neg(fgraph: FunctionGraph, node: Node) -> list[TensorVar
     if isinstance(node.op, MeasurableMaxNeg):
         return None  # pragma: no cover
 
-    base_var = node.inputs[0]
+    base_var = cast(TensorVariable, node.inputs[0])
 
     # Min is the Max of the negation of the same distribution. Hence, op must be Elemwise
     if not (base_var.owner is not None and isinstance(base_var.owner.op, Elemwise)):
@@ -213,6 +216,7 @@ def find_measurable_max_neg(fgraph: FunctionGraph, node: Node) -> list[TensorVar
         return None
 
     # distinguish measurable discrete and continuous (because logprob is different)
+    measurable_min: Max
     if base_rv.owner.op.dtype.startswith("int"):
         measurable_min = MeasurableDiscreteMaxNeg(list(axis))
     else:

pymc/pytensorf.py

@@ -926,15 +926,15 @@ def find_default_update(clients, rng: Variable) -> None | Variable:
     if inputs is None:
         inputs = []
 
-    outputs = makeiter(outputs)
-    fg = FunctionGraph(outputs=outputs, clone=False)
+    outs = makeiter(outputs)
+    fg = FunctionGraph(outputs=outs, clone=False)
     clients = fg.clients
 
     rng_updates = {}
     # Iterate over input RNGs. Only consider shared RNGs if `must_be_shared==True`
     for input_rng in (
         inp
-        for inp in graph_inputs(outputs, blockers=inputs)
+        for inp in graph_inputs(outs, blockers=inputs)
         if (
             (not must_be_shared or isinstance(inp, SharedVariable))
             and isinstance(inp.type, RandomType)

pymc/smc/sampling.py

@@ -225,7 +225,7 @@ def sample_smc(
     trace = MultiTrace(traces)
 
     _t_sampling = time.time() - t1
-    sample_stats, idata = _save_sample_stats(
+    _, idata = _save_sample_stats(
         sample_settings,
         sample_stats,
         chains,

scripts/run_mypy.py

@@ -25,7 +25,6 @@
 pymc/distributions/continuous.py
 pymc/distributions/dist_math.py
 pymc/distributions/distribution.py
-pymc/distributions/mixture.py
 pymc/distributions/multivariate.py
 pymc/distributions/timeseries.py
 pymc/distributions/truncated.py
@@ -34,7 +33,6 @@
 pymc/logprob/censoring.py
 pymc/logprob/basic.py
 pymc/logprob/mixture.py
-pymc/logprob/order.py
 pymc/logprob/rewriting.py
 pymc/logprob/scan.py
 pymc/logprob/tensor.py