Add Coords and StrongCoords typing aliases and standardize model/arviz usage

Author: aman-coder03

pymc/backends/arviz.py

@@ -23,8 +23,12 @@
     Optional,
     Union,
     cast,
+    TypeAlias,
 )
 
+from pymc.distributions.shape_utils import StrongCoords
+
+
 import numpy as np
 import xarray
 
@@ -56,6 +60,7 @@
 
 # random variable object ...
 Var = Any
+DimsDict: TypeAlias = Mapping[str, Sequence[str]]
 
 
 def dict_to_dataset_drop_incompatible_coords(vars_dict, *args, dims, coords, **kwargs):
@@ -123,7 +128,8 @@ def find_constants(model: "Model") -> dict[str, Var]:
     return constant_data
 
 
-def coords_and_dims_for_inferencedata(model: Model) -> tuple[dict[str, Any], dict[str, Any]]:
+def coords_and_dims_for_inferencedata(model: Model,) -> tuple[StrongCoords, DimsDict]:
+
     """Parse PyMC model coords and dims format to one accepted by InferenceData."""
     coords = {
         cname: np.array(cvals) if isinstance(cvals, tuple) else cvals

pymc/distributions/shape_utils.py

@@ -97,6 +97,15 @@ def _check_shape_type(shape):
 StrongDimsWithEllipsis: TypeAlias = Sequence[str | EllipsisType]
 StrongSize: TypeAlias = TensorVariable | tuple[int | Variable, ...]
 
+from collections.abc import Mapping
+from typing import Hashable
+
+CoordValue: TypeAlias = Sequence[Hashable] | np.ndarray | None
+Coords: TypeAlias = Mapping[str, CoordValue]
+
+StrongCoordValue: TypeAlias = tuple[Hashable, ...] | None
+StrongCoords: TypeAlias = Mapping[str, StrongCoordValue]
+
 
 def convert_dims(dims: Dims | None) -> StrongDims | None:
     """Process a user-provided dims variable into None or a valid dims tuple."""

pymc/model/core.py

@@ -20,6 +20,8 @@
 import warnings
 
 from collections.abc import Iterable, Sequence
+from pymc.distributions.shape_utils import Coords, StrongCoords, CoordValue
+
 from typing import (
     Literal,
     cast,
@@ -453,7 +455,7 @@ def _validate_name(name):
     def __init__(
         self,
         name="",
-        coords=None,
+        coords: Coords | None = None,
         check_bounds=True,
         *,
         model: _UnsetType | None | Model = UNSET,
@@ -488,7 +490,7 @@ def __init__(
             self.deterministics = treelist()
             self.potentials = treelist()
             self.data_vars = treelist()
-            self._coords = {}
+            self._coords: StrongCoords = {}
             self._dim_lengths = {}
         self.add_coords(coords)
 
@@ -907,9 +909,9 @@ def unobserved_RVs(self):
         return self.free_RVs + self.deterministics
 
     @property
-    def coords(self) -> dict[str, tuple | None]:
+    def coords(self) -> StrongCoords:
         """Coordinate values for model dimensions."""
-        return self._coords
+        return self._coords 
 
     @property
     def dim_lengths(self) -> dict[str, TensorVariable]:
@@ -937,7 +939,7 @@ def shape_from_dims(self, dims):
     def add_coord(
         self,
         name: str,
-        values: Sequence | np.ndarray | None = None,
+        values: CoordValue = None,
         *,
         length: int | Variable | None = None,
     ):