Use new-style typing annotations, not comments

Author: Armavica

pymc/backends/arviz.py

@@ -26,13 +26,11 @@
 import pymc
 
 from pymc.aesaraf import extract_obs_data
-from pymc.model import modelcontext
+from pymc.model import Model, modelcontext
 from pymc.util import get_default_varnames
 
 if TYPE_CHECKING:
-
     from pymc.backends.base import MultiTrace  # pylint: disable=invalid-name
-    from pymc.model import Model
 
 ___all__ = [""]
 
@@ -144,12 +142,10 @@ def insert(self, k: str, v, idx: int):
 class InferenceDataConverter:  # pylint: disable=too-many-instance-attributes
     """Encapsulate InferenceData specific logic."""
 
-    model = None  # type: Optional[Model]
-    nchains = None  # type: int
-    ndraws = None  # type: int
-    posterior_predictive = None  # Type: Optional[Mapping[str, np.ndarray]]
-    predictions = None  # Type: Optional[Mapping[str, np.ndarray]]
-    prior = None  # Type: Optional[Mapping[str, np.ndarray]]
+    model: Optional[Model] = None
+    posterior_predictive: Optional[Mapping[str, np.ndarray]] = None
+    predictions: Optional[Mapping[str, np.ndarray]] = None
+    prior: Optional[Mapping[str, np.ndarray]] = None
 
     def __init__(
         self,

pymc/backends/ndarray.py

@@ -85,7 +85,7 @@ def setup(self, draws, chain, sampler_vars=None) -> None:
         if self._stats is None:
             self._stats = []
             for sampler in sampler_vars:
-                data = dict()  # type: Dict[str, np.ndarray]
+                data: Dict[str, np.ndarray] = dict()
                 self._stats.append(data)
                 for varname, dtype in sampler.items():
                     data[varname] = np.zeros(draws, dtype=dtype)

pymc/data.py

@@ -301,7 +301,7 @@ class Minibatch(TensorVariable):
         >>> assert x.eval().shape == (2, 20, 20, 40, 10)
     """
 
-    RNG = collections.defaultdict(list)  # type: Dict[str, List[Any]]
+    RNG: Dict[str, List[Any]] = collections.defaultdict(list)
 
     @aesara.config.change_flags(compute_test_value="raise")
     def __init__(

pymc/distributions/distribution.py

@@ -63,9 +63,9 @@
 
 DIST_PARAMETER_TYPES: TypeAlias = Union[np.ndarray, int, float, TensorVariable]
 
-vectorized_ppc = contextvars.ContextVar(
+vectorized_ppc: contextvars.ContextVar[Optional[Callable]] = contextvars.ContextVar(
     "vectorized_ppc", default=None
-)  # type: contextvars.ContextVar[Optional[Callable]]
+)
 
 PLATFORM = sys.platform
 

pymc/model.py

@@ -188,7 +188,7 @@ def get_context(cls, error_if_none=True) -> Optional[T]:
         on the stack, or ``None``. If ``error_if_none`` is True (default),
         raise a ``TypeError`` instead of returning ``None``."""
         try:
-            candidate = cls.get_contexts()[-1]  # type: Optional[T]
+            candidate: Optional[T] = cls.get_contexts()[-1]
         except IndexError as e:
             # Calling code expects to get a TypeError if the entity
             # is unfound, and there's too much to fix.

pymc/variational/opvi.py

@@ -1482,7 +1482,7 @@ def sample(
 
         if random_seed is not None:
             (random_seed,) = _get_seeds_per_chain(random_seed, 1)
-        samples = self.sample_dict_fn(draws, random_seed=random_seed)  # type: dict
+        samples: dict = self.sample_dict_fn(draws, random_seed=random_seed)
         points = ({name: records[i] for name, records in samples.items()} for i in range(draws))
 
         trace = NDArray(