A few typing fixes

Author: Armavica

pymc/distributions/dist_math.py

@@ -61,10 +61,10 @@ def check_parameters(logp: Variable, *conditions: Iterable[Variable], msg: str =
     check_bounds = False in pm.Model()
     """
     # at.all does not accept True/False, but accepts np.array(True)/np.array(False)
-    conditions = [
+    conditions_ = [
         cond if (cond is not True and cond is not False) else np.array(cond) for cond in conditions
     ]
-    all_true_scalar = at.all([at.all(cond) for cond in conditions])
+    all_true_scalar = at.all([at.all(cond) for cond in conditions_])
     return CheckParameterValue(msg)(logp, all_true_scalar)
 
 

pymc/model_graph.py

@@ -73,7 +73,7 @@ def _expand(x):
             return []
 
         parents = {
-            get_var_name(x)
+            VarName(get_var_name(x))
             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
@@ -109,7 +109,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 [var.name for var in selected_ancestors]
+        return [VarName(var.name) for var in selected_ancestors]
 
     def make_compute_graph(
         self, var_names: Optional[Iterable[VarName]] = None
@@ -230,7 +230,7 @@ def get_plates(self, var_names: Optional[Iterable[VarName]] = None) -> Dict[str,
                 plate_label = " x ".join(dim_labels)
             else:
                 # The RV has no `dims` information.
-                dim_labels = map(str, shape)
+                dim_labels = [str(x) for x in shape]
                 plate_label = " x ".join(map(str, shape))
             plates[plate_label].add(var_name)
 

pymc/sampling/jax.py

@@ -48,8 +48,8 @@
     get_default_varnames,
 )
 
-xla_flags = os.getenv("XLA_FLAGS", "")
-xla_flags = re.sub(r"--xla_force_host_platform_device_count=.+\s", "", xla_flags).split()
+xla_flags_env = os.getenv("XLA_FLAGS", "")
+xla_flags = re.sub(r"--xla_force_host_platform_device_count=.+\s", "", xla_flags_env).split()
 os.environ["XLA_FLAGS"] = " ".join([f"--xla_force_host_platform_device_count={100}"] + xla_flags)
 
 __all__ = (
@@ -108,7 +108,7 @@ def get_jaxified_graph(
 ) -> List[TensorVariable]:
     """Compile an PyTensor graph into an optimized JAX function"""
 
-    graph = _replace_shared_variables(outputs)
+    graph = _replace_shared_variables(outputs) if outputs is not None else None
 
     fgraph = FunctionGraph(inputs=inputs, outputs=graph, clone=True)
     # We need to add a Supervisor to the fgraph to be able to run the
@@ -251,12 +251,10 @@ def _get_batched_jittered_initial_points(
         jitter=jitter,
         jitter_max_retries=jitter_max_retries,
     )
-    initial_points = [list(initial_point.values()) for initial_point in initial_points]
+    initial_points_values = [list(initial_point.values()) for initial_point in initial_points]
     if chains == 1:
-        initial_points = initial_points[0]
-    else:
-        initial_points = [np.stack(init_state) for init_state in zip(*initial_points)]
-    return initial_points
+        return initial_points_values[0]
+    return [np.stack(init_state) for init_state in zip(*initial_points_values)]
 
 
 def _update_coords_and_dims(

pymc/variational/opvi.py

@@ -51,6 +51,8 @@
 import itertools
 import warnings
 
+from typing import Any
+
 import numpy as np
 import pytensor
 import pytensor.tensor as at
@@ -673,11 +675,11 @@ class Group(WithMemoization):
     initial_dist_map = 0.0
 
     # for handy access using class methods
-    __param_spec__ = dict()
+    __param_spec__: dict = dict()
     short_name = ""
-    alias_names = frozenset()
-    __param_registry = dict()
-    __name_registry = dict()
+    alias_names: frozenset[str] = frozenset()
+    __param_registry: dict[frozenset, Any] = dict()
+    __name_registry: dict[str, Any] = dict()
 
     @classmethod
     def register(cls, sbcls):
@@ -1552,11 +1554,11 @@ def sample(
         finally:
             trace.close()
 
-        trace = MultiTrace([trace])
+        multi_trace = MultiTrace([trace])
         if not return_inferencedata:
-            return trace
+            return multi_trace
         else:
-            return pm.to_inference_data(trace, model=self.model, **kwargs)
+            return pm.to_inference_data(multi_trace, model=self.model, **kwargs)
 
     @property
     def ndim(self):