Move coords typing to pymc.typing and fix circular imports

Author: aman-coder03

pymc/backends/arviz.py

@@ -39,14 +39,16 @@
 
 import pymc
 
-from pymc.distributions.shape_utils import StrongCoords
-from pymc.model import Model, modelcontext
-from pymc.progress_bar import CustomProgress, default_progress_theme
-from pymc.pytensorf import PointFunc, extract_obs_data
-from pymc.util import get_default_varnames
+from pymc.model import modelcontext
+from pymc.typing import StrongCoords
 
 if TYPE_CHECKING:
     from pymc.backends.base import MultiTrace
+    from pymc.model import Model
+
+from pymc.progress_bar import CustomProgress, default_progress_theme
+from pymc.pytensorf import PointFunc, extract_obs_data
+from pymc.util import get_default_varnames
 
 ___all__ = [""]
 

pymc/distributions/shape_utils.py

@@ -16,7 +16,7 @@
 
 import warnings
 
-from collections.abc import Hashable, Mapping, Sequence
+from collections.abc import Sequence
 from functools import singledispatch
 from types import EllipsisType
 from typing import Any, TypeAlias, cast
@@ -97,12 +97,6 @@ def _check_shape_type(shape):
 StrongDimsWithEllipsis: TypeAlias = Sequence[str | EllipsisType]
 StrongSize: TypeAlias = TensorVariable | tuple[int | Variable, ...]
 
-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."""
@@ -416,6 +410,7 @@ def get_support_shape(
         if len(dims) < ndim_supp:
             raise ValueError(f"Number of dims is too small for ndim_supp of {ndim_supp}")
         from pymc.model.core import modelcontext
+
         model = modelcontext(None)
         inferred_support_shape = [
             model.dim_lengths[dims[i]] - support_shape_offset[i] for i in range(-ndim_supp, 0)

pymc/step_methods/state.py

@@ -30,7 +30,7 @@ class DataClassState:
 def equal_dataclass_values(v1, v2):
     if v1.__class__ != v2.__class__:
         return False
-    if isinstance(v1, (list, tuple)):  # noqa: UP038
+    if isinstance(v1, (list, tuple)):
         return len(v1) == len(v2) and all(
             equal_dataclass_values(v1i, v2i) for v1i, v2i in zip(v1, v2, strict=True)
         )

pymc/typing.py

@@ -0,0 +1,33 @@
+#   Copyright 2024 - present The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+
+from __future__ import annotations
+
+from collections.abc import Hashable, Mapping, Sequence
+from typing import TypeAlias
+
+import numpy as np
+
+# -------------------------
+# Coordinate typing helpers
+# -------------------------
+
+# User-facing coordinate values (before normalization)
+CoordValue: TypeAlias = Sequence[Hashable] | np.ndarray | None
+Coords: TypeAlias = Mapping[str, CoordValue]
+
+# After normalization / internal representation
+StrongCoordValue: TypeAlias = tuple[Hashable, ...] | None
+StrongCoords: TypeAlias = Mapping[str, StrongCoordValue]