From 9baf9ca5d6b08ee4b4f52635c85698b90c95a842 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Tue, 2 Dec 2025 18:31:48 +0530 Subject: [PATCH 01/26] Add Coords and StrongCoords typing aliases and standardize model/arviz usage --- pymc/backends/arviz.py | 8 +++++++- pymc/distributions/shape_utils.py | 9 +++++++++ pymc/model/core.py | 12 +++++++----- 3 files changed, 23 insertions(+), 6 deletions(-) diff --git a/pymc/backends/arviz.py b/pymc/backends/arviz.py index 63f8370523..9910271f76 100644 --- a/pymc/backends/arviz.py +++ b/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 diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index cdab3046b1..0123e3e664 100644 --- a/pymc/distributions/shape_utils.py +++ b/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.""" diff --git a/pymc/model/core.py b/pymc/model/core.py index 3630138e00..8a407e5ed2 100644 --- a/pymc/model/core.py +++ b/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, ): From 7fc13121fbcf2f627778d8c71c8212af54a56e83 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 3 Dec 2025 22:37:22 +0530 Subject: [PATCH 02/26] Fix ruff formatting, imports, and dev requirements --- pymc/backends/arviz.py | 11 +++++------ pymc/distributions/shape_utils.py | 5 +---- pymc/model/core.py | 5 ++--- 3 files changed, 8 insertions(+), 13 deletions(-) diff --git a/pymc/backends/arviz.py b/pymc/backends/arviz.py index 9910271f76..dba26582f3 100644 --- a/pymc/backends/arviz.py +++ b/pymc/backends/arviz.py @@ -21,14 +21,11 @@ TYPE_CHECKING, Any, Optional, + TypeAlias, Union, cast, - TypeAlias, ) -from pymc.distributions.shape_utils import StrongCoords - - import numpy as np import xarray @@ -42,6 +39,7 @@ 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 @@ -128,8 +126,9 @@ def find_constants(model: "Model") -> dict[str, Var]: return constant_data -def coords_and_dims_for_inferencedata(model: Model,) -> tuple[StrongCoords, DimsDict]: - +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 diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 0123e3e664..11c19de559 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -16,7 +16,7 @@ import warnings -from collections.abc import Sequence +from collections.abc import Hashable, Mapping, Sequence from functools import singledispatch from types import EllipsisType from typing import Any, TypeAlias, cast @@ -97,9 +97,6 @@ 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] diff --git a/pymc/model/core.py b/pymc/model/core.py index 8a407e5ed2..7a58e52f66 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -20,8 +20,6 @@ import warnings from collections.abc import Iterable, Sequence -from pymc.distributions.shape_utils import Coords, StrongCoords, CoordValue - from typing import ( Literal, cast, @@ -46,6 +44,7 @@ from pymc.blocking import DictToArrayBijection, RaveledVars from pymc.data import MinibatchOp, is_valid_observed +from pymc.distributions.shape_utils import Coords, CoordValue, StrongCoords from pymc.exceptions import ( BlockModelAccessError, ImputationWarning, @@ -911,7 +910,7 @@ def unobserved_RVs(self): @property def coords(self) -> StrongCoords: """Coordinate values for model dimensions.""" - return self._coords + return self._coords @property def dim_lengths(self) -> dict[str, TensorVariable]: From 9e60df8191bb4fa08f7b58fa8f0033374f1ef948 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 3 Dec 2025 22:47:09 +0530 Subject: [PATCH 03/26] Fix circular import by importing modelcontext from pymc.model.core --- pymc/distributions/shape_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 11c19de559..e0832c3cbe 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -33,7 +33,7 @@ from pytensor.tensor.type_other import NoneTypeT from pytensor.tensor.variable import TensorVariable -from pymc.model import modelcontext +from pymc.model.core import modelcontext from pymc.pytensorf import convert_observed_data __all__ = [ From 46ed9cade3d6e3ee4bf937efe38d97cf82ee42d3 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 3 Dec 2025 22:53:51 +0530 Subject: [PATCH 04/26] Fix circular import by lazily importing modelcontext in shape_from_dims --- pymc/distributions/shape_utils.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index e0832c3cbe..cbd8935037 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -33,7 +33,7 @@ from pytensor.tensor.type_other import NoneTypeT from pytensor.tensor.variable import TensorVariable -from pymc.model.core import modelcontext +from pymc.model import modelcontext from pymc.pytensorf import convert_observed_data __all__ = [ @@ -170,21 +170,27 @@ def convert_size(size: Size) -> StrongSize | None: ) -def shape_from_dims(dims: StrongDims, model) -> StrongShape: +def shape_from_dims(dims: StrongDims, model=None) -> StrongShape: """Determine shape from a `dims` tuple. Parameters ---------- dims : array-like A vector of dimension names or None. - model : pm.Model - The current model on stack. + model : pm.Model, optional + The current model on stack. If None, it will be resolved via modelcontext. Returns ------- - dims : tuple of (str or None) - Names or None for all RV dimensions. + shape : tuple + Shape inferred from model dimension lengths. """ + # Lazy import to break circular dependency + if model is None: + from pymc.model.core import modelcontext + + model = modelcontext(None) + # Dims must be known already unknowndim_dims = set(dims) - set(model.dim_lengths) if unknowndim_dims: From 8f5040d01ad88bdbacf12e00c266ab16cfc982a6 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 3 Dec 2025 22:58:56 +0530 Subject: [PATCH 05/26] Fix circular import by using only lazy modelcontext imports --- pymc/distributions/shape_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index cbd8935037..7a61aece33 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -33,7 +33,7 @@ from pytensor.tensor.type_other import NoneTypeT from pytensor.tensor.variable import TensorVariable -from pymc.model import modelcontext +#from pymc.model import modelcontext from pymc.pytensorf import convert_observed_data __all__ = [ From 02f9053ad17bd0e3b9b56c7c7b45dfae716b08e6 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 3 Dec 2025 23:09:02 +0530 Subject: [PATCH 06/26] Fix Model circular import using TYPE_CHECKING and lazy import --- pymc/printing.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pymc/printing.py b/pymc/printing.py index 63514ac4d0..e769aeb831 100644 --- a/pymc/printing.py +++ b/pymc/printing.py @@ -14,8 +14,8 @@ import re - from functools import partial +from typing import TYPE_CHECKING from pytensor.compile import SharedVariable from pytensor.graph.basic import Constant @@ -26,7 +26,9 @@ from pytensor.tensor.random.type import RandomType from pytensor.tensor.type_other import NoneTypeT -from pymc.model import Model +if TYPE_CHECKING: + from pymc.model import Model + __all__ = [ "str_for_dist", @@ -301,6 +303,7 @@ def _default_repr_pretty(obj: TensorVariable | Model, p, cycle): try: # register our custom pretty printer in ipython shells import IPython + from pymc.model.core import Model IPython.lib.pretty.for_type(TensorVariable, _default_repr_pretty) IPython.lib.pretty.for_type(Model, _default_repr_pretty) From c6f3f34b07889b30214953aeb324837366a45827 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 3 Dec 2025 23:12:33 +0530 Subject: [PATCH 07/26] Fix lazy modelcontext import flagged by ruff --- pymc/distributions/shape_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 7a61aece33..cdc51b47f2 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -33,7 +33,7 @@ from pytensor.tensor.type_other import NoneTypeT from pytensor.tensor.variable import TensorVariable -#from pymc.model import modelcontext +# from pymc.model import modelcontext from pymc.pytensorf import convert_observed_data __all__ = [ @@ -415,6 +415,8 @@ def get_support_shape( assert isinstance(dims, tuple) 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) From 5b4b1f6c5996c517a38f362e7cb761231792f21b Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 3 Dec 2025 23:19:32 +0530 Subject: [PATCH 08/26] Fix missing modelcontext import flagged by ruff --- pymc/distributions/shape_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index cdc51b47f2..46132c23a7 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -416,7 +416,6 @@ 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) From ba05761a301e3676b039d314ce20a844e8d369f2 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 3 Dec 2025 23:25:39 +0530 Subject: [PATCH 09/26] Fix circular import of Model in printing.py --- pymc/printing.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pymc/printing.py b/pymc/printing.py index e769aeb831..8796dc8981 100644 --- a/pymc/printing.py +++ b/pymc/printing.py @@ -14,6 +14,7 @@ import re + from functools import partial from typing import TYPE_CHECKING @@ -303,6 +304,7 @@ def _default_repr_pretty(obj: TensorVariable | Model, p, cycle): try: # register our custom pretty printer in ipython shells import IPython + from pymc.model.core import Model IPython.lib.pretty.for_type(TensorVariable, _default_repr_pretty) From 21b04fdbfb005c9765af4c5b859a5a9319a3a850 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Thu, 4 Dec 2025 12:35:19 +0530 Subject: [PATCH 10/26] Move coords typing to pymc.typing and fix circular imports --- pymc/backends/arviz.py | 12 ++++++----- pymc/distributions/shape_utils.py | 9 ++------- pymc/step_methods/state.py | 2 +- pymc/typing.py | 33 +++++++++++++++++++++++++++++++ 4 files changed, 43 insertions(+), 13 deletions(-) create mode 100644 pymc/typing.py diff --git a/pymc/backends/arviz.py b/pymc/backends/arviz.py index dba26582f3..42d4e401e2 100644 --- a/pymc/backends/arviz.py +++ b/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__ = [""] diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 46132c23a7..91f49497ff 100644 --- a/pymc/distributions/shape_utils.py +++ b/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) diff --git a/pymc/step_methods/state.py b/pymc/step_methods/state.py index 98e177aa03..3cd765624d 100644 --- a/pymc/step_methods/state.py +++ b/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) ) diff --git a/pymc/typing.py b/pymc/typing.py new file mode 100644 index 0000000000..577675b964 --- /dev/null +++ b/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] From b5ca2083226ae8417b337e0e4465152fc1ed9905 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Thu, 4 Dec 2025 12:47:03 +0530 Subject: [PATCH 11/26] Fix ruff UP038 isinstance union style --- pymc/step_methods/state.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/step_methods/state.py b/pymc/step_methods/state.py index 3cd765624d..6fe97de21a 100644 --- a/pymc/step_methods/state.py +++ b/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)): + 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) ) From 0f9cad7c867bc6d94e9e95eaafdd15f8fac8ecf4 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Thu, 4 Dec 2025 12:56:17 +0530 Subject: [PATCH 12/26] Fix printing Model NameError and move Coord typing to pymc.typing --- pymc/model/core.py | 2 +- pymc/printing.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 7a58e52f66..6c12b3460e 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -44,7 +44,7 @@ from pymc.blocking import DictToArrayBijection, RaveledVars from pymc.data import MinibatchOp, is_valid_observed -from pymc.distributions.shape_utils import Coords, CoordValue, StrongCoords +from pymc.typing import Coords, CoordValue, StrongCoords from pymc.exceptions import ( BlockModelAccessError, ImputationWarning, diff --git a/pymc/printing.py b/pymc/printing.py index 8796dc8981..180038e650 100644 --- a/pymc/printing.py +++ b/pymc/printing.py @@ -13,6 +13,8 @@ # limitations under the License. +from __future__ import annotations + import re from functools import partial From 26d756d36b7604553677d5db3140df628a9e185f Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Thu, 4 Dec 2025 13:09:09 +0530 Subject: [PATCH 13/26] Move coords typing to pymc.typing and fix printing imports --- pymc/backends/arviz.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/pymc/backends/arviz.py b/pymc/backends/arviz.py index 42d4e401e2..54edd3c6bf 100644 --- a/pymc/backends/arviz.py +++ b/pymc/backends/arviz.py @@ -13,6 +13,8 @@ # limitations under the License. """PyMC-ArviZ conversion code.""" +from __future__ import annotations + import logging import warnings @@ -20,9 +22,7 @@ from typing import ( TYPE_CHECKING, Any, - Optional, TypeAlias, - Union, cast, ) @@ -90,7 +90,7 @@ def dict_to_dataset_drop_incompatible_coords(vars_dict, *args, dims, coords, **k return dict_to_dataset(vars_dict, *args, dims=dims, coords=safe_coords, **kwargs) -def find_observations(model: "Model") -> dict[str, Var]: +def find_observations(model: Model) -> dict[str, Var]: """If there are observations available, return them as a dictionary.""" observations = {} for obs in model.observed_RVs: @@ -107,7 +107,7 @@ def find_observations(model: "Model") -> dict[str, Var]: return observations -def find_constants(model: "Model") -> dict[str, Var]: +def find_constants(model: Model) -> dict[str, Var]: """If there are constants available, return them as a dictionary.""" model_vars = model.basic_RVs + model.deterministics + model.potentials value_vars = set(model.rvs_to_values.values()) @@ -272,7 +272,7 @@ def __init__( self.observations = find_observations(self.model) - def split_trace(self) -> tuple[Union[None, "MultiTrace"], Union[None, "MultiTrace"]]: + def split_trace(self) -> tuple[None | MultiTrace, None | MultiTrace]: """Split MultiTrace object into posterior and warmup. Returns @@ -498,7 +498,7 @@ def to_inference_data(self): def to_inference_data( - trace: Optional["MultiTrace"] = None, + trace: MultiTrace | None = None, *, prior: Mapping[str, Any] | None = None, posterior_predictive: Mapping[str, Any] | None = None, @@ -507,7 +507,7 @@ def to_inference_data( coords: CoordSpec | None = None, dims: DimSpec | None = None, sample_dims: list | None = None, - model: Optional["Model"] = None, + model: Model | None = None, save_warmup: bool | None = None, include_transformed: bool = False, ) -> InferenceData: @@ -575,8 +575,8 @@ def to_inference_data( ### perhaps we should have an inplace argument? def predictions_to_inference_data( predictions, - posterior_trace: Optional["MultiTrace"] = None, - model: Optional["Model"] = None, + posterior_trace: MultiTrace | None = None, + model: Model | None = None, coords: CoordSpec | None = None, dims: DimSpec | None = None, sample_dims: list | None = None, From 447f990dd661cf6753d3115a3e7da2ce03c77d15 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Thu, 4 Dec 2025 13:15:09 +0530 Subject: [PATCH 14/26] Remove implicit modelcontext fallback from shape_from_dims --- pymc/distributions/shape_utils.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 91f49497ff..1b74f06627 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -164,26 +164,23 @@ def convert_size(size: Size) -> StrongSize | None: ) -def shape_from_dims(dims: StrongDims, model=None) -> StrongShape: +def shape_from_dims(dims: StrongDims, model: Model) -> StrongShape: """Determine shape from a `dims` tuple. Parameters ---------- dims : array-like A vector of dimension names or None. - model : pm.Model, optional - The current model on stack. If None, it will be resolved via modelcontext. + model : pm.Model + The current model on stack. Returns ------- shape : tuple Shape inferred from model dimension lengths. """ - # Lazy import to break circular dependency if model is None: - from pymc.model.core import modelcontext - - model = modelcontext(None) + raise ValueError("model must be provided explicitly to infer shape from dims") # Dims must be known already unknowndim_dims = set(dims) - set(model.dim_lengths) From 627e4030d5b2154adb9e1895766406b8e908ea77 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Thu, 4 Dec 2025 13:25:44 +0530 Subject: [PATCH 15/26] Fix shape_from_dims typing and remove circular import --- pymc/distributions/shape_utils.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 1b74f06627..98355267b1 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -19,7 +19,7 @@ from collections.abc import Sequence from functools import singledispatch from types import EllipsisType -from typing import Any, TypeAlias, cast +from typing import TYPE_CHECKING, Any, TypeAlias, cast import numpy as np @@ -33,9 +33,12 @@ from pytensor.tensor.type_other import NoneTypeT from pytensor.tensor.variable import TensorVariable -# from pymc.model import modelcontext from pymc.pytensorf import convert_observed_data +if TYPE_CHECKING: + from pymc.model import Model + + __all__ = [ "change_dist_size", "rv_size_is_none", @@ -164,7 +167,7 @@ def convert_size(size: Size) -> StrongSize | None: ) -def shape_from_dims(dims: StrongDims, model: Model) -> StrongShape: +def shape_from_dims(dims: StrongDims, model: "Model") -> StrongShape: """Determine shape from a `dims` tuple. Parameters From 784de8b596f5b65c1afb34720f822d5c58680166 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Thu, 4 Dec 2025 13:27:55 +0530 Subject: [PATCH 16/26] Fix typing and import order --- pymc/model/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/model/core.py b/pymc/model/core.py index 6c12b3460e..fbc0ca0fb6 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -44,7 +44,6 @@ from pymc.blocking import DictToArrayBijection, RaveledVars from pymc.data import MinibatchOp, is_valid_observed -from pymc.typing import Coords, CoordValue, StrongCoords from pymc.exceptions import ( BlockModelAccessError, ImputationWarning, @@ -66,6 +65,7 @@ join_nonshared_inputs, rewrite_pregrad, ) +from pymc.typing import Coords, CoordValue, StrongCoords from pymc.util import ( UNSET, WithMemoization, From dec6f7c13a6c9564123c5899077c5c98b85cc627 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Sat, 6 Dec 2025 19:14:38 +0530 Subject: [PATCH 17/26] Removing comments from typing.py --- pymc/typing.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/pymc/typing.py b/pymc/typing.py index 577675b964..a1f779a09b 100644 --- a/pymc/typing.py +++ b/pymc/typing.py @@ -20,10 +20,6 @@ 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] From 087ebe56723ec353ad9947c42e34450bc25f46e1 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Sat, 6 Dec 2025 19:32:20 +0530 Subject: [PATCH 18/26] Remove deprecated pymc.typing module after moving aliases to pymc.util --- pymc/backends/arviz.py | 2 +- pymc/model/core.py | 5 ++++- pymc/typing.py | 29 ----------------------------- pymc/util.py | 12 +++++++++++- 4 files changed, 16 insertions(+), 32 deletions(-) delete mode 100644 pymc/typing.py diff --git a/pymc/backends/arviz.py b/pymc/backends/arviz.py index 54edd3c6bf..272f5f755f 100644 --- a/pymc/backends/arviz.py +++ b/pymc/backends/arviz.py @@ -40,7 +40,7 @@ import pymc from pymc.model import modelcontext -from pymc.typing import StrongCoords +from pymc.util import StrongCoords if TYPE_CHECKING: from pymc.backends.base import MultiTrace diff --git a/pymc/model/core.py b/pymc/model/core.py index fbc0ca0fb6..884332bba1 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -65,8 +65,10 @@ join_nonshared_inputs, rewrite_pregrad, ) -from pymc.typing import Coords, CoordValue, StrongCoords from pymc.util import ( + Coords, + CoordValue, + StrongCoords, UNSET, WithMemoization, _UnsetType, @@ -76,6 +78,7 @@ treedict, treelist, ) + from pymc.vartypes import continuous_types, discrete_types, typefilter __all__ = [ diff --git a/pymc/typing.py b/pymc/typing.py deleted file mode 100644 index a1f779a09b..0000000000 --- a/pymc/typing.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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 - -# 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] diff --git a/pymc/util.py b/pymc/util.py index 32d8d65e70..6c18ad8585 100644 --- a/pymc/util.py +++ b/pymc/util.py @@ -18,7 +18,7 @@ from collections import namedtuple from collections.abc import Sequence from copy import deepcopy -from typing import cast +from typing import Mapping, TypeAlias, Hashable, cast import arviz import cloudpickle @@ -31,6 +31,16 @@ from pymc.exceptions import BlockModelAccessError +#Coordinate & Shape Typing +CoordValue: TypeAlias = Sequence[Hashable] | np.ndarray | None +Coords: TypeAlias = Mapping[str, CoordValue] + +StrongCoordValue: TypeAlias = tuple[Hashable, ...] | None +StrongCoords: TypeAlias = Mapping[str, StrongCoordValue] + +StrongDims: TypeAlias = tuple[str, ...] +StrongShape: TypeAlias = tuple[int, ...] + class _UnsetType: """Type for the `UNSET` object to make it look nice in `help(...)` outputs.""" From 0016b6bd9322ebbb485714a134a4c30cd2b47aa7 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Sat, 6 Dec 2025 19:47:39 +0530 Subject: [PATCH 19/26] Fixes --- pymc/backends/arviz.py | 2 +- pymc/distributions/shape_utils.py | 37 ++++++++++++++----------------- pymc/model/core.py | 3 +-- pymc/util.py | 28 +++++++++++++++++++---- 4 files changed, 43 insertions(+), 27 deletions(-) diff --git a/pymc/backends/arviz.py b/pymc/backends/arviz.py index 272f5f755f..98b8598eb0 100644 --- a/pymc/backends/arviz.py +++ b/pymc/backends/arviz.py @@ -40,7 +40,7 @@ import pymc from pymc.model import modelcontext -from pymc.util import StrongCoords +from pymc.util import StrongCoords if TYPE_CHECKING: from pymc.backends.base import MultiTrace diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 98355267b1..244822e9e9 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -6,7 +6,7 @@ # # http://www.apache.org/licenses/LICENSE-2.0 # -# Unless required by applicable law or agreed to in writing, software +# Unless required by applicable law or deemed 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 @@ -14,12 +14,14 @@ """Common shape operations to broadcast samples from probability distributions for stochastic nodes in PyMC.""" +from __future__ import annotations + import warnings from collections.abc import Sequence from functools import singledispatch from types import EllipsisType -from typing import TYPE_CHECKING, Any, TypeAlias, cast +from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -33,11 +35,22 @@ from pytensor.tensor.type_other import NoneTypeT from pytensor.tensor.variable import TensorVariable -from pymc.pytensorf import convert_observed_data +from pymc.exceptions import ShapeError +from pymc.pytensorf import PotentialShapeType, convert_observed_data +from pymc.util import StrongDims, StrongShape if TYPE_CHECKING: from pymc.model import Model +Shape = int | TensorVariable | Sequence[int | Variable] +Dims = str | Sequence[str | None] +DimsWithEllipsis = str | EllipsisType | Sequence[str | None | EllipsisType] +Size = int | TensorVariable | Sequence[int | Variable] +# Strong (validated) types from util + +# Additional strong types needed inside this file +StrongDimsWithEllipsis = Sequence[str | EllipsisType] +StrongSize = TensorVariable | tuple[int | Variable, ...] __all__ = [ "change_dist_size", @@ -45,9 +58,6 @@ "to_tuple", ] -from pymc.exceptions import ShapeError -from pymc.pytensorf import PotentialShapeType - def to_tuple(shape): """Convert ints, arrays, and Nones to tuples. @@ -88,19 +98,6 @@ def _check_shape_type(shape): return tuple(out) -# User-provided can be lazily specified as scalars -Shape: TypeAlias = int | TensorVariable | Sequence[int | Variable] -Dims: TypeAlias = str | Sequence[str | None] -DimsWithEllipsis: TypeAlias = str | EllipsisType | Sequence[str | None | EllipsisType] -Size: TypeAlias = int | TensorVariable | Sequence[int | Variable] - -# After conversion to vectors -StrongShape: TypeAlias = TensorVariable | tuple[int | Variable, ...] -StrongDims: TypeAlias = Sequence[str] -StrongDimsWithEllipsis: TypeAlias = Sequence[str | EllipsisType] -StrongSize: TypeAlias = TensorVariable | tuple[int | Variable, ...] - - def convert_dims(dims: Dims | None) -> StrongDims | None: """Process a user-provided dims variable into None or a valid dims tuple.""" if dims is None: @@ -167,7 +164,7 @@ def convert_size(size: Size) -> StrongSize | None: ) -def shape_from_dims(dims: StrongDims, model: "Model") -> StrongShape: +def shape_from_dims(dims: StrongDims, model: Model) -> StrongShape: """Determine shape from a `dims` tuple. Parameters diff --git a/pymc/model/core.py b/pymc/model/core.py index 884332bba1..20c2e3ec47 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -66,10 +66,10 @@ rewrite_pregrad, ) from pymc.util import ( + UNSET, Coords, CoordValue, StrongCoords, - UNSET, WithMemoization, _UnsetType, get_transformed_name, @@ -78,7 +78,6 @@ treedict, treelist, ) - from pymc.vartypes import continuous_types, discrete_types, typefilter __all__ = [ diff --git a/pymc/util.py b/pymc/util.py index 6c18ad8585..c9d6943e09 100644 --- a/pymc/util.py +++ b/pymc/util.py @@ -16,9 +16,10 @@ import re from collections import namedtuple -from collections.abc import Sequence +from collections.abc import Hashable, Mapping, Sequence from copy import deepcopy -from typing import Mapping, TypeAlias, Hashable, cast +from types import EllipsisType +from typing import TypeAlias, cast import arviz import cloudpickle @@ -26,21 +27,40 @@ import xarray from cachetools import LRUCache, cachedmethod -from pytensor import Variable from pytensor.compile import SharedVariable +from pytensor.graph.basic import Variable +from pytensor.tensor.variable import TensorVariable from pymc.exceptions import BlockModelAccessError -#Coordinate & Shape Typing +# ---- User-facing coordinate types ---- CoordValue: TypeAlias = Sequence[Hashable] | np.ndarray | None Coords: TypeAlias = Mapping[str, CoordValue] +# ---- Internal strong coordinate types ---- StrongCoordValue: TypeAlias = tuple[Hashable, ...] | None StrongCoords: TypeAlias = Mapping[str, StrongCoordValue] +# ---- Internal strong dimension/shape types ---- StrongDims: TypeAlias = tuple[str, ...] StrongShape: TypeAlias = tuple[int, ...] +# User-provided shape before processing +Shape: TypeAlias = int | TensorVariable | Sequence[int | Variable] + +# User-provided dims before processing +Dims: TypeAlias = str | Sequence[str | None] + +# User-provided dims that may include ellipsis (...) +DimsWithEllipsis: TypeAlias = str | EllipsisType | Sequence[str | None | EllipsisType] + +# User-provided size before processing +Size: TypeAlias = int | TensorVariable | Sequence[int | Variable] + +# Strong / normalized versions used internally +StrongDimsWithEllipsis: TypeAlias = Sequence[str | EllipsisType] +StrongSize: TypeAlias = TensorVariable | tuple[int | Variable, ...] + class _UnsetType: """Type for the `UNSET` object to make it look nice in `help(...)` outputs.""" From c930d5b50bc17930a31e9f5faf08e7796407da06 Mon Sep 17 00:00:00 2001 From: Aman Srivastava <160766756+aman-coder03@users.noreply.github.com> Date: Tue, 9 Dec 2025 19:18:40 +0530 Subject: [PATCH 20/26] Apply suggestions from code review Co-authored-by: Michael Osthege --- pymc/distributions/shape_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 244822e9e9..469d8e2e9e 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -6,7 +6,7 @@ # # http://www.apache.org/licenses/LICENSE-2.0 # -# Unless required by applicable law or deemed in writing, software +# 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 From c651cb2eccb2f89b74b02ba30cbde3a2c5c1562e Mon Sep 17 00:00:00 2001 From: Aman Srivastava <160766756+aman-coder03@users.noreply.github.com> Date: Tue, 9 Dec 2025 19:23:31 +0530 Subject: [PATCH 21/26] Update util.py --- pymc/util.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/pymc/util.py b/pymc/util.py index c9d6943e09..2c7a42ee12 100644 --- a/pymc/util.py +++ b/pymc/util.py @@ -33,15 +33,15 @@ from pymc.exceptions import BlockModelAccessError -# ---- User-facing coordinate types ---- +# User-facing coordinate types CoordValue: TypeAlias = Sequence[Hashable] | np.ndarray | None Coords: TypeAlias = Mapping[str, CoordValue] -# ---- Internal strong coordinate types ---- +# Internal strong coordinate types StrongCoordValue: TypeAlias = tuple[Hashable, ...] | None StrongCoords: TypeAlias = Mapping[str, StrongCoordValue] -# ---- Internal strong dimension/shape types ---- +# Internal strong dimension/shape types StrongDims: TypeAlias = tuple[str, ...] StrongShape: TypeAlias = tuple[int, ...] @@ -60,6 +60,17 @@ # Strong / normalized versions used internally StrongDimsWithEllipsis: TypeAlias = Sequence[str | EllipsisType] StrongSize: TypeAlias = TensorVariable | tuple[int | Variable, ...] +""" +Type alias groups: + +- User-facing types (Coords, CoordValue, Shape, Dims, Size) represent the flexible inputs + accepted from users when defining models. + +- Strong/internal types (StrongCoords, StrongDims, StrongShape, StrongSize, etc.) represent + normalized, fully-validated forms used internally after processing. + +These distinctions allow looser user input while keeping strict internal consistency. +""" class _UnsetType: From fb07885de4a15e0d1836560c4acfd376eb428b46 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Tue, 9 Dec 2025 19:37:05 +0530 Subject: [PATCH 22/26] Refactor shape_from_dims usage to Model method --- pymc/distributions/distribution.py | 3 +-- pymc/distributions/shape_utils.py | 30 +----------------------------- 2 files changed, 2 insertions(+), 31 deletions(-) diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 27d53c8687..35a929f550 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -48,7 +48,6 @@ convert_size, find_size, rv_size_is_none, - shape_from_dims, ) from pymc.logprob.abstract import MeasurableOp, _icdf, _logcdf, _logprob from pymc.logprob.basic import logp @@ -522,7 +521,7 @@ def __new__( # finally, observed, to determine the shape of the variable. if kwargs.get("size") is None and kwargs.get("shape") is None: if dims is not None: - kwargs["shape"] = shape_from_dims(dims, model) + kwargs["shape"] = model.shape_from_dims(dims) elif observed is not None: kwargs["shape"] = tuple(observed.shape) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 469d8e2e9e..28b93ee89e 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -40,7 +40,7 @@ from pymc.util import StrongDims, StrongShape if TYPE_CHECKING: - from pymc.model import Model + pass Shape = int | TensorVariable | Sequence[int | Variable] Dims = str | Sequence[str | None] DimsWithEllipsis = str | EllipsisType | Sequence[str | None | EllipsisType] @@ -164,34 +164,6 @@ def convert_size(size: Size) -> StrongSize | None: ) -def shape_from_dims(dims: StrongDims, model: Model) -> StrongShape: - """Determine shape from a `dims` tuple. - - Parameters - ---------- - dims : array-like - A vector of dimension names or None. - model : pm.Model - The current model on stack. - - Returns - ------- - shape : tuple - Shape inferred from model dimension lengths. - """ - if model is None: - raise ValueError("model must be provided explicitly to infer shape from dims") - - # Dims must be known already - unknowndim_dims = set(dims) - set(model.dim_lengths) - if unknowndim_dims: - raise KeyError( - f"Dimensions {unknowndim_dims} are unknown to the model and cannot be used to specify a `shape`." - ) - - return tuple(model.dim_lengths[dname] for dname in dims) - - def find_size( shape: StrongShape | None, size: StrongSize | None, From 7c22c16388eefa03e7bdf63188c0b56c153143de Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 10 Dec 2025 22:08:56 +0530 Subject: [PATCH 23/26] implementing the suggested changes --- pymc/distributions/distribution.py | 2 +- pymc/model/core.py | 2 +- pymc/util.py | 35 ++++++++++++++---------------- 3 files changed, 18 insertions(+), 21 deletions(-) diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 35a929f550..edc315bd3e 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -521,7 +521,7 @@ def __new__( # finally, observed, to determine the shape of the variable. if kwargs.get("size") is None and kwargs.get("shape") is None: if dims is not None: - kwargs["shape"] = model.shape_from_dims(dims) + kwargs["shape"] = model.symbolic_shape_from_dims(dims) elif observed is not None: kwargs["shape"] = tuple(observed.shape) diff --git a/pymc/model/core.py b/pymc/model/core.py index 20c2e3ec47..2a093baffe 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -922,7 +922,7 @@ def dim_lengths(self) -> dict[str, TensorVariable]: """ return self._dim_lengths - def shape_from_dims(self, dims): + def symbolic_shape_from_dims(self, dims): shape = [] if len(set(dims)) != len(dims): raise ValueError("Can not contain the same dimension name twice.") diff --git a/pymc/util.py b/pymc/util.py index 2c7a42ee12..c9a6dec4fc 100644 --- a/pymc/util.py +++ b/pymc/util.py @@ -33,44 +33,41 @@ from pymc.exceptions import BlockModelAccessError -# User-facing coordinate types CoordValue: TypeAlias = Sequence[Hashable] | np.ndarray | None +# User-provided values for a single coordinate dimension. + Coords: TypeAlias = Mapping[str, CoordValue] +# Mapping from dimension name to its coordinate values. -# Internal strong coordinate types StrongCoordValue: TypeAlias = tuple[Hashable, ...] | None +# Normalized coordinate values stored internally. + StrongCoords: TypeAlias = Mapping[str, StrongCoordValue] +# Mapping from dimension name to normalized coordinate values. -# Internal strong dimension/shape types StrongDims: TypeAlias = tuple[str, ...] +# Tuple of dimension names after validation. + StrongShape: TypeAlias = tuple[int, ...] +# Fully-resolved numeric shape used internally. -# User-provided shape before processing Shape: TypeAlias = int | TensorVariable | Sequence[int | Variable] +# User-provided shape specification before normalization. -# User-provided dims before processing Dims: TypeAlias = str | Sequence[str | None] +# User-provided dimension names before normalization. -# User-provided dims that may include ellipsis (...) DimsWithEllipsis: TypeAlias = str | EllipsisType | Sequence[str | None | EllipsisType] +# User-provided dimension names that may include ellipsis. -# User-provided size before processing Size: TypeAlias = int | TensorVariable | Sequence[int | Variable] +# User-provided size specification before normalization. -# Strong / normalized versions used internally StrongDimsWithEllipsis: TypeAlias = Sequence[str | EllipsisType] -StrongSize: TypeAlias = TensorVariable | tuple[int | Variable, ...] -""" -Type alias groups: +# Normalized dimension names that may include ellipsis. -- User-facing types (Coords, CoordValue, Shape, Dims, Size) represent the flexible inputs - accepted from users when defining models. - -- Strong/internal types (StrongCoords, StrongDims, StrongShape, StrongSize, etc.) represent - normalized, fully-validated forms used internally after processing. - -These distinctions allow looser user input while keeping strict internal consistency. -""" +StrongSize: TypeAlias = TensorVariable | tuple[int | Variable, ...] +# Normalized symbolic size used internally. class _UnsetType: From e823b9fa77de4ec707d43e45e581a2bb6cf303fc Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 17 Dec 2025 01:05:03 +0530 Subject: [PATCH 24/26] Use shared type aliases from pymc.util in shape_utils --- pymc/distributions/shape_utils.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 28b93ee89e..abdda0edd1 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -37,20 +37,17 @@ from pymc.exceptions import ShapeError from pymc.pytensorf import PotentialShapeType, convert_observed_data -from pymc.util import StrongDims, StrongShape +from pymc.util import ( + Shape, + Dims, + DimsWithEllipsis, + Size, + StrongDims, + StrongShape, + StrongDimsWithEllipsis, + StrongSize, +) -if TYPE_CHECKING: - pass -Shape = int | TensorVariable | Sequence[int | Variable] -Dims = str | Sequence[str | None] -DimsWithEllipsis = str | EllipsisType | Sequence[str | None | EllipsisType] -Size = int | TensorVariable | Sequence[int | Variable] - -# Strong (validated) types from util - -# Additional strong types needed inside this file -StrongDimsWithEllipsis = Sequence[str | EllipsisType] -StrongSize = TensorVariable | tuple[int | Variable, ...] __all__ = [ "change_dist_size", From 4f924a1b556562869f166511d49abeca9ebdc857 Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 17 Dec 2025 01:11:15 +0530 Subject: [PATCH 25/26] Fix pre-commit formatting issues --- pymc/distributions/shape_utils.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index abdda0edd1..ed0550bb31 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -21,7 +21,7 @@ from collections.abc import Sequence from functools import singledispatch from types import EllipsisType -from typing import TYPE_CHECKING, Any, cast +from typing import Any, cast import numpy as np @@ -38,17 +38,16 @@ from pymc.exceptions import ShapeError from pymc.pytensorf import PotentialShapeType, convert_observed_data from pymc.util import ( - Shape, Dims, DimsWithEllipsis, + Shape, Size, StrongDims, - StrongShape, StrongDimsWithEllipsis, + StrongShape, StrongSize, ) - __all__ = [ "change_dist_size", "rv_size_is_none", From 14f77d6b95ce724737e6e77db2d5edeb91a42c1a Mon Sep 17 00:00:00 2001 From: Aman Srivastava Date: Wed, 17 Dec 2025 02:19:06 +0530 Subject: [PATCH 26/26] Fix docstring --- pymc/util.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/pymc/util.py b/pymc/util.py index c9a6dec4fc..32d39efac9 100644 --- a/pymc/util.py +++ b/pymc/util.py @@ -34,40 +34,40 @@ from pymc.exceptions import BlockModelAccessError CoordValue: TypeAlias = Sequence[Hashable] | np.ndarray | None -# User-provided values for a single coordinate dimension. +"""User-provided values for a single coordinate dimension.""" Coords: TypeAlias = Mapping[str, CoordValue] -# Mapping from dimension name to its coordinate values. +"""Mapping from dimension name to its coordinate values.""" StrongCoordValue: TypeAlias = tuple[Hashable, ...] | None -# Normalized coordinate values stored internally. +"""Normalized coordinate values stored internally.""" StrongCoords: TypeAlias = Mapping[str, StrongCoordValue] -# Mapping from dimension name to normalized coordinate values. +"""Mapping from dimension name to normalized coordinate values.""" StrongDims: TypeAlias = tuple[str, ...] -# Tuple of dimension names after validation. +"""Tuple of dimension names after validation.""" StrongShape: TypeAlias = tuple[int, ...] -# Fully-resolved numeric shape used internally. +"""Fully-resolved numeric shape used internally.""" Shape: TypeAlias = int | TensorVariable | Sequence[int | Variable] -# User-provided shape specification before normalization. +"""User-provided shape specification before normalization.""" Dims: TypeAlias = str | Sequence[str | None] -# User-provided dimension names before normalization. +"""User-provided dimension names before normalization.""" DimsWithEllipsis: TypeAlias = str | EllipsisType | Sequence[str | None | EllipsisType] -# User-provided dimension names that may include ellipsis. +"""User-provided dimension names that may include ellipsis.""" Size: TypeAlias = int | TensorVariable | Sequence[int | Variable] -# User-provided size specification before normalization. +"""User-provided size specification before normalization.""" StrongDimsWithEllipsis: TypeAlias = Sequence[str | EllipsisType] -# Normalized dimension names that may include ellipsis. +"""Normalized dimension names that may include ellipsis.""" StrongSize: TypeAlias = TensorVariable | tuple[int | Variable, ...] -# Normalized symbolic size used internally. +"""Normalized symbolic size used internally.""" class _UnsetType: