Fixes

Author: aman-coder03

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

pymc/distributions/shape_utils.py

@@ -6,20 +6,22 @@
 #
 #       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
 #   limitations under the License.
 
 """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,21 +35,29 @@
 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",
     "rv_size_is_none",
     "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

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__ = [

pymc/util.py

@@ -16,31 +16,51 @@
 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
 import numpy as np
 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."""