Fix typing in multiple places

Author: michaelosthege

pymc/printing.py

@@ -13,6 +13,8 @@
 #   limitations under the License.
 
 
+from functools import partial
+
 from pytensor.compile import SharedVariable
 from pytensor.graph.basic import Constant, walk
 from pytensor.tensor.basic import TensorVariable, Variable
@@ -55,7 +57,7 @@ def str_for_dist(
 
     if "latex" in formatting:
         if print_name is not None:
-            print_name = r"\text{" + _latex_escape(dist.name.strip("$")) + "}"
+            print_name = r"\text{" + _latex_escape(print_name.strip("$")) + "}"
 
         op_name = (
             dist.owner.op._print_name[1]
@@ -96,17 +98,16 @@ def str_for_model(model: Model, formatting: str = "plain", include_params: bool
     """Make a human-readable string representation of Model, listing all random variables
     and their distributions, optionally including parameter values."""
 
-    kwargs = dict(formatting=formatting, include_params=include_params)
-    free_rv_reprs = [str_for_dist(dist, **kwargs) for dist in model.free_RVs]
-    observed_rv_reprs = [str_for_dist(rv, **kwargs) for rv in model.observed_RVs]
-    det_reprs = [
-        str_for_potential_or_deterministic(dist, **kwargs, dist_name="Deterministic")
-        for dist in model.deterministics
-    ]
-    potential_reprs = [
-        str_for_potential_or_deterministic(pot, **kwargs, dist_name="Potential")
-        for pot in model.potentials
-    ]
+    # Wrap functions to avoid confusing typecheckers
+    sfd = partial(str_for_dist, formatting=formatting, include_params=include_params)
+    sfp = partial(
+        str_for_potential_or_deterministic, formatting=formatting, include_params=include_params
+    )
+
+    free_rv_reprs = [sfd(dist) for dist in model.free_RVs]
+    observed_rv_reprs = [sfd(rv) for rv in model.observed_RVs]
+    det_reprs = [sfp(dist, dist_name="Deterministic") for dist in model.deterministics]
+    potential_reprs = [sfp(pot, dist_name="Potential") for pot in model.potentials]
 
     var_reprs = free_rv_reprs + det_reprs + observed_rv_reprs + potential_reprs
 
@@ -162,6 +163,8 @@ def _str_for_input_var(var: Variable, formatting: str) -> str:
     from pymc.distributions.distribution import SymbolicRandomVariable
 
     def _is_potential_or_deterministic(var: Variable) -> bool:
+        if not hasattr(var, "str_repr"):
+            return False
         try:
             return var.str_repr.__func__.func is str_for_potential_or_deterministic
         except AttributeError:
@@ -175,14 +178,15 @@ def _is_potential_or_deterministic(var: Variable) -> bool:
     ) or _is_potential_or_deterministic(var):
         # show the names for RandomVariables, Deterministics, and Potentials, rather
         # than the full expression
+        assert isinstance(var, TensorVariable)
         return _str_for_input_rv(var, formatting)
     elif isinstance(var.owner.op, DimShuffle):
         return _str_for_input_var(var.owner.inputs[0], formatting)
     else:
         return _str_for_expression(var, formatting)
 
 
-def _str_for_input_rv(var: Variable, formatting: str) -> str:
+def _str_for_input_rv(var: TensorVariable, formatting: str) -> str:
     _str = (
         var.name
         if var.name is not None
@@ -221,12 +225,15 @@ def _expand(x):
         if x.owner and (not isinstance(x.owner.op, RandomVariable | SymbolicRandomVariable)):
             return reversed(x.owner.inputs)
 
-    parents = [
-        x
-        for x in walk(nodes=var.owner.inputs, expand=_expand)
-        if x.owner and isinstance(x.owner.op, RandomVariable | SymbolicRandomVariable)
-    ]
-    names = [x.name for x in parents]
+    parents = []
+    names = []
+    for x in walk(nodes=var.owner.inputs, expand=_expand):
+        assert isinstance(x, Variable)
+        if x.owner and isinstance(x.owner.op, RandomVariable | SymbolicRandomVariable):
+            parents.append(x)
+            xname = x.name
+            assert xname is not None
+            names.append(xname)
 
     if "latex" in formatting:
         return (
@@ -257,6 +264,8 @@ def _default_repr_pretty(obj: TensorVariable | Model, p, cycle):
     """Handy plug-in method to instruct IPython-like REPLs to use our str_repr above."""
     # we know that our str_repr does not recurse, so we can ignore cycle
     try:
+        if not hasattr(obj, "str_repr"):
+            raise AttributeError
         output = obj.str_repr()
         # Find newlines and replace them with p.break_()
         # (see IPython.lib.pretty._repr_pprint)

pymc/pytensorf.py

@@ -14,6 +14,7 @@
 import warnings
 
 from collections.abc import Callable, Generator, Iterable, Sequence
+from typing import cast
 
 import numpy as np
 import pandas as pd
@@ -29,7 +30,6 @@
 from pytensor.graph.basic import (
     Apply,
     Constant,
-    Node,
     Variable,
     clone_get_equiv,
     graph_inputs,
@@ -208,8 +208,8 @@ def replace_vars_in_graphs(
     """
     # Clone graphs and get equivalences
     inputs = [i for i in graph_inputs(graphs) if not isinstance(i, Constant)]
-    equiv = {k: k for k in replacements.keys()}
-    equiv = clone_get_equiv(inputs, graphs, False, False, equiv)
+    memo = {k: k for k in replacements.keys()}
+    equiv = clone_get_equiv(inputs, graphs, False, False, memo)
 
     fg = FunctionGraph(
         [equiv[i] for i in inputs],
@@ -753,7 +753,7 @@ def find_rng_nodes(
     ]
 
 
-def replace_rng_nodes(outputs: Sequence[TensorVariable]) -> Sequence[TensorVariable]:
+def replace_rng_nodes(outputs: Sequence[TensorVariable]) -> list[TensorVariable]:
     """Replace any RNG nodes upstream of outputs by new RNGs of the same type
 
     This can be used when combining a pre-existing graph with a cloned one, to ensure
@@ -775,7 +775,7 @@ def replace_rng_nodes(outputs: Sequence[TensorVariable]) -> Sequence[TensorVaria
             rng_cls = np.random.Generator
         new_rng_nodes.append(pytensor.shared(rng_cls(np.random.PCG64())))
     graph.replace_all(zip(rng_nodes, new_rng_nodes), import_missing=True)
-    return graph.outputs
+    return cast(list[TensorVariable], graph.outputs)
 
 
 SeedSequenceSeed = None | int | Sequence[int] | np.ndarray | np.random.SeedSequence
@@ -798,7 +798,7 @@ def reseed_rngs(
         rng.set_value(new_rng, borrow=True)
 
 
-def collect_default_updates_inner_fgraph(node: Node) -> dict[Variable, Variable]:
+def collect_default_updates_inner_fgraph(node: Apply) -> dict[Variable, Variable]:
     """Collect default updates from node with inner fgraph."""
     op = node.op
     inner_updates = collect_default_updates(
@@ -945,7 +945,7 @@ def find_default_update(clients, rng: Variable) -> None | Variable:
         default_update = find_default_update(clients, input_rng)
 
         # Respect default update if provided
-        if getattr(input_rng, "default_update", None):
+        if hasattr(input_rng, "default_update") and input_rng.default_update is not None:
             rng_updates[input_rng] = input_rng.default_update
         else:
             if default_update is not None:
@@ -1001,7 +1001,8 @@ def compile_pymc(
 
     # We always reseed random variables as this provides RNGs with no chances of collision
     if rng_updates:
-        reseed_rngs(rng_updates.keys(), random_seed)
+        rngs = cast(list[SharedVariable], list(rng_updates))
+        reseed_rngs(rngs, random_seed)
 
     # If called inside a model context, see if check_bounds flag is set to False
     try:

scripts/run_mypy.py

@@ -44,7 +44,6 @@
 pymc/model/core.py
 pymc/model/fgraph.py
 pymc/model/transform/conditioning.py
-pymc/printing.py
 pymc/pytensorf.py
 pymc/sampling/jax.py
 """