Skip to content

Conversation

@aman-coder03
Copy link

@aman-coder03 aman-coder03 commented Dec 3, 2025

this PR addresses issue #7972 by consolidating the typing of coords across the PyMC codebase.

Previously, coords appeared with a mix of types such as:

  • Sequence
  • Sequence | np.ndarray
  • dict[str, Any]

this PR introduces explicit and consistent type aliases:

  • CoordValue
  • Coords
  • StrongCoordValue
  • StrongCoords

(similar to the existing Shape, Dims, and StrongDims aliases in shape_utils)

Key changes

  • standardizes Model.coords to return StrongCoords
  • updates ArviZ conversion to return StrongCoords instead of dict[str, Any}
  • removes the mixed use of Sequence, Sequence | np.ndarray, and untyped dictionaries
  • improves static typing and downstream library compatibility

Notes

while testing this change, a runtime circular import involving:
pymc.model.core
pymc.distributions.shape_utils
from pymc.model import modelcontext
pymc.model (partially initialized)
was encountered.

Fixes #7972

Comment on lines 188 to 192
# Lazy import to break circular dependency
if model is None:
from pymc.model.core import modelcontext

model = modelcontext(None)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why this functional change? Doesn't seem necessary or a good idea

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i introduced the lazy import only to work around the circular import that surfaced during the typing consolidation (shape_utils -> model -> distributions -> shape_utils). I agree that changing runtime behavior here is not ideal
if you prefer, i can revert this functional change and instead resolve the circular dependency purely by relocating the typing aliases to a more independent module

Comment on lines 100 to 104
CoordValue: TypeAlias = Sequence[Hashable] | np.ndarray | None
Coords: TypeAlias = Mapping[str, CoordValue]

StrongCoordValue: TypeAlias = tuple[Hashable, ...] | None
StrongCoords: TypeAlias = Mapping[str, StrongCoordValue]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Put these in a more independent place to avoid the new circular import issues. Coordinates are not really distribution specific anyway

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i can move CoordValue, Coords, StrongCoordValue, and StrongCoords to a more independent location and update downstream imports accordingly. Let me know if you have a preferred target module for these.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Consolidate typing for coords

2 participants