Refactor shape_from_dims usage to Model method

Author: aman-coder03

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)
 

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,