.Remove with dims (merge this before)

Author: ricardoV94

pymc/dims/__init__.py

@@ -68,4 +68,4 @@ def __init__():
 
 from pymc.dims import math
 from pymc.dims.distributions import *
-from pymc.dims.model import Data, with_dims
+from pymc.dims.model import Data, Deterministic, Potential

pymc/dims/distribution_core.py

@@ -24,7 +24,6 @@
 from pytensor.xtensor.type import XTensorVariable
 
 from pymc import SymbolicRandomVariable, modelcontext
-from pymc.dims.model import with_dims
 from pymc.dims.transforms import DimTransform, log_odds_transform, log_transform
 from pymc.distributions.distribution import _support_point, support_point
 from pymc.distributions.shape_utils import DimsWithEllipsis, convert_dims
@@ -129,14 +128,11 @@ def _as_xtensor(x):
         try:
             return as_xtensor(x)
         except TypeError:
-            try:
-                return with_dims(x)
-            except ValueError:
-                raise ValueError(
-                    f"Variable {x} must have dims associated with it.\n"
-                    "To avoid subtle bugs, PyMC does not make any assumptions about the dims of the parameters.\n"
-                    "Convert parameters to an xarray.DataArray, pymc.dims.Data or pytensor.xtensor.as_xtensor with explicit dims."
-                )
+            raise ValueError(
+                f"Variable {x} must have dims associated with it.\n"
+                "To avoid subtle bugs, PyMC does not make any assumptions about the dims of the parameters.\n"
+                "Convert parameters to an xarray.DataArray, pymc.dims.Data or pytensor.xtensor.as_xtensor with explicit dims."
+            )
 
     def __new__(
         cls,

pymc/dims/model.py

@@ -15,7 +15,6 @@
 
 from pytensor.tensor import TensorVariable
 from pytensor.xtensor import as_xtensor
-from pytensor.xtensor.basic import TensorFromXTensor
 from pytensor.xtensor.type import XTensorVariable
 
 from pymc.data import Data as RegularData
@@ -25,29 +24,6 @@
 from pymc.model.core import Potential as RegularPotential
 
 
-def with_dims(x: TensorVariable | XTensorVariable, model: Model | None = None) -> XTensorVariable:
-    """Recover the dims of a variable that was registered in the Model."""
-    if isinstance(x, XTensorVariable):
-        return x
-
-    if (x.owner is not None) and isinstance(x.owner.op, TensorFromXTensor):
-        dims = x.owner.inputs[0].type.dims
-        return as_xtensor(x, dims=dims, name=x.name)
-
-    # Try accessing the model context to get dims
-    try:
-        model = modelcontext(model)
-        if (
-            model.named_vars.get(x.name, None) is x
-            and (dims := model.named_vars_to_dims.get(x.name, None)) is not None
-        ):
-            return as_xtensor(x, dims=dims, name=x.name)
-    except TypeError:
-        pass
-
-    raise ValueError(f"variable {x} doesn't have dims associated with it")
-
-
 def Data(
     name: str, value, dims: DimsWithEllipsis = None, model: Model | None = None, **kwargs
 ) -> XTensorVariable: