More precise type-hints

Author: ricardoV94

pymc/distributions/shape_utils.py

@@ -369,7 +369,7 @@ def get_support_shape(
         support_shape_offset = [0] * ndim_supp
     elif isinstance(support_shape_offset, int):
         support_shape_offset = [support_shape_offset] * ndim_supp
-    inferred_support_shape: Sequence[int | np.ndarray | Variable] | None = None
+    inferred_support_shape: Sequence[int | np.ndarray | TensorVariable] | None = None
 
     if shape is not None:
         shape = to_tuple(shape)
@@ -411,7 +411,7 @@ def get_support_shape(
         # There were two sources of support_shape, make sure they are consistent
         inferred_support_shape = [
             cast(
-                Variable,
+                TensorVariable,
                 Assert(msg="support_shape does not match respective shape dimension")(
                     inferred, pt.eq(inferred, explicit)
                 ),

pymc/gp/cov.py

@@ -164,6 +164,7 @@ class Covariance(BaseCovariance):
 
     def __init__(self, input_dim: int, active_dims: IntSequence | None = None):
         self.input_dim = input_dim
+        self.active_dims: np.ndarray[Any, np.dtype[np.int_]]
         if active_dims is None:
             self.active_dims = np.arange(input_dim)
         else:

pymc/model/core.py

@@ -895,7 +895,7 @@ def coords(self) -> dict[str, tuple | None]:
         return self._coords
 
     @property
-    def dim_lengths(self) -> dict[str, Variable]:
+    def dim_lengths(self) -> dict[str, TensorVariable]:
         """The symbolic lengths of dimensions in the model.
 
         The values are typically instances of ``TensorVariable`` or ``ScalarSharedVariable``.