Fix typing in logprob.tensor

Author: michaelosthege

pymc/logprob/rewriting.py

@@ -36,7 +36,7 @@
 import warnings
 
 from collections import deque
-from collections.abc import Sequence
+from collections.abc import Collection, Sequence
 
 import pytensor.tensor as pt
 
@@ -484,7 +484,7 @@ def cleanup_ir(vars: Sequence[Variable]) -> None:
 
 
 def assume_measured_ir_outputs(
-    inputs: Sequence[TensorVariable], outputs: Sequence[TensorVariable]
+    inputs: Collection[TensorVariable], outputs: Sequence[TensorVariable]
 ) -> Sequence[TensorVariable]:
     """Run IR rewrite assuming each output is measured.
 

pymc/logprob/tensor.py

@@ -38,6 +38,7 @@
 import pytensor
 
 from pytensor import tensor as pt
+from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import compute_test_value
 from pytensor.graph.rewriting.basic import node_rewriter
 from pytensor.tensor import TensorVariable
@@ -60,7 +61,7 @@
 
 
 @node_rewriter([Alloc])
-def naive_bcast_rv_lift(fgraph, node):
+def naive_bcast_rv_lift(fgraph: FunctionGraph, node):
     """Lift an ``Alloc`` through a ``RandomVariable`` ``Op``.
 
     XXX: This implementation simply broadcasts the ``RandomVariable``'s
@@ -226,6 +227,7 @@ def find_measurable_stacks(fgraph, node) -> list[TensorVariable] | None:
         measurable_stack = MeasurableJoin()(axis, *base_vars)
     else:
         measurable_stack = MeasurableMakeVector(node.op.dtype)(*base_vars)
+    assert isinstance(measurable_stack, TensorVariable)
 
     return [measurable_stack]
 
@@ -242,7 +244,7 @@ class MeasurableDimShuffle(DimShuffle):
 
 
 @_logprob.register(MeasurableDimShuffle)
-def logprob_dimshuffle(op, values, base_var, **kwargs):
+def logprob_dimshuffle(op: MeasurableDimShuffle, values, base_var, **kwargs):
     """Compute the log-likelihood graph for a `MeasurableDimShuffle`."""
     (value,) = values
 
@@ -300,6 +302,7 @@ def find_measurable_dimshuffles(fgraph, node) -> list[TensorVariable] | None:
     measurable_dimshuffle = MeasurableDimShuffle(node.op.input_broadcastable, node.op.new_order)(
         base_var
     )
+    assert isinstance(measurable_dimshuffle, TensorVariable)
 
     return [measurable_dimshuffle]
 

pymc/logprob/utils.py

@@ -36,7 +36,7 @@
 import typing
 import warnings
 
-from collections.abc import Container, Sequence
+from collections.abc import Container, Iterable, Sequence
 
 import numpy as np
 import pytensor
@@ -173,7 +173,7 @@ def indices_from_subtensor(idx_list, indices):
 
 
 def check_potential_measurability(
-    inputs: tuple[TensorVariable], valued_rvs: Container[TensorVariable]
+    inputs: Iterable[TensorVariable], valued_rvs: Container[TensorVariable]
 ) -> bool:
     valued_rvs = set(valued_rvs)
 

scripts/run_mypy.py

@@ -36,7 +36,6 @@
 pymc/logprob/mixture.py
 pymc/logprob/rewriting.py
 pymc/logprob/scan.py
-pymc/logprob/tensor.py
 pymc/logprob/transform_value.py
 pymc/logprob/transforms.py
 pymc/logprob/utils.py