Fix typing in opvi and model_graph

Author: michaelosthege

pymc/model_graph.py

@@ -76,12 +76,12 @@ def _expand(x):
                 return reversed(_filter_non_parameter_inputs(x))
             return []
 
-        parents = {
-            VarName(get_var_name(x))
-            for x in walk(nodes=_filter_non_parameter_inputs(var), expand=_expand)
+        parents = set()
+        for x in walk(nodes=_filter_non_parameter_inputs(var), expand=_expand):
             # Only consider nodes that are in the named model variables.
-            if x.name and x.name in self._all_var_names
-        }
+            vname = getattr(x, "name", None)
+            if isinstance(vname, str) and vname in self._all_var_names:
+                parents.add(VarName(vname))
 
         return parents
 
@@ -113,7 +113,7 @@ def vars_to_plot(self, var_names: Optional[Iterable[VarName]] = None) -> List[Va
                 selected_ancestors.add(self.model.rvs_to_values[var])
 
         # ordering of self._all_var_names is important
-        return [VarName(var.name) for var in selected_ancestors]
+        return [VarName(get_var_name(var)) for var in selected_ancestors]
 
     def make_compute_graph(
         self, var_names: Optional[Iterable[VarName]] = None

pymc/variational/opvi.py

@@ -51,7 +51,7 @@
 import itertools
 import warnings
 
-from typing import Any
+from typing import Any, overload
 
 import numpy as np
 import pytensor
@@ -980,17 +980,29 @@ def symbolic_random(self):
         """
         raise NotImplementedError
 
-    @pytensor.config.change_flags(compute_test_value="off")
+    @overload
+    def set_size_and_deterministic(
+        self, node: Variable, s, d: bool, more_replacements: dict | None = None
+    ) -> Variable:
+        ...
+
+    @overload
     def set_size_and_deterministic(
-        self, node: Variable, s, d: bool, more_replacements: dict = None
+        self, node: list[Variable], s, d: bool, more_replacements: dict | None = None
     ) -> list[Variable]:
+        ...
+
+    @pytensor.config.change_flags(compute_test_value="off")
+    def set_size_and_deterministic(
+        self, node: Variable | list[Variable], s, d: bool, more_replacements: dict | None = None
+    ) -> Variable | list[Variable]:
         """*Dev* - after node is sampled via :func:`symbolic_sample_over_posterior` or
         :func:`symbolic_single_sample` new random generator can be allocated and applied to node
 
         Parameters
         ----------
-        node: :class:`Variable`
-            PyTensor node with symbolically applied VI replacements
+        node
+            PyTensor node(s) with symbolically applied VI replacements
         s: scalar
             desired number of samples
         d: bool or int
@@ -1000,7 +1012,7 @@ def set_size_and_deterministic(
 
         Returns
         -------
-        :class:`Variable` with applied replacements, ready to use
+        :class:`Variable` or list with applied replacements, ready to use
         """
 
         flat2rand = self.make_size_and_deterministic_replacements(s, d, more_replacements)

scripts/run_mypy.py

@@ -43,11 +43,9 @@
 pymc/model/fgraph.py
 pymc/model/transform/basic.py
 pymc/model/transform/conditioning.py
-pymc/model_graph.py
 pymc/printing.py
 pymc/pytensorf.py
 pymc/sampling/jax.py
-pymc/variational/opvi.py
 """