Implement xarray like semantics in dims module

Author: ricardoV94

.github/workflows/tests.yml

@@ -135,6 +135,12 @@ jobs:
             tests/logprob/test_transforms.py
             tests/logprob/test_utils.py
 
+          - |
+            tests/dims/distributions/test_core.py
+            tests/dims/distributions/test_scalar.py
+            tests/dims/distributions/test_vector.py
+            tests/dims/test_model.py
+
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:

conda-envs/environment-alternative-backends.yml

@@ -22,7 +22,7 @@ dependencies:
 - numpyro>=0.8.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.5,<2.32
 - python-graphviz
 - networkx
 - rich>=13.7.1

conda-envs/environment-dev.yml

@@ -12,7 +12,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.5,<2.32
 - python-graphviz
 - networkx
 - scipy>=1.4.1

conda-envs/environment-docs.yml

@@ -11,7 +11,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.5,<2.32
 - python-graphviz
 - rich>=13.7.1
 - scipy>=1.4.1

conda-envs/environment-test.yml

@@ -14,7 +14,7 @@ dependencies:
 - pandas>=0.24.0
 - pip
 - polyagamma
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.5,<2.32
 - python-graphviz
 - networkx
 - rich>=13.7.1

conda-envs/windows-environment-dev.yml

@@ -12,7 +12,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.5,<2.32
 - python-graphviz
 - networkx
 - rich>=13.7.1

conda-envs/windows-environment-test.yml

@@ -15,7 +15,7 @@ dependencies:
 - pandas>=0.24.0
 - pip
 - polyagamma
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.5,<2.32
 - python-graphviz
 - networkx
 - rich>=13.7.1

pymc/data.py

@@ -13,11 +13,12 @@
 #   limitations under the License.
 
 import io
+import typing
 import urllib.request
 
 from collections.abc import Sequence
 from copy import copy
-from typing import cast
+from typing import Union, cast
 
 import numpy as np
 import pandas as pd
@@ -32,12 +33,13 @@
 from pytensor.tensor.random.basic import IntegersRV
 from pytensor.tensor.variable import TensorConstant, TensorVariable
 
-import pymc as pm
-
-from pymc.logprob.utils import rvs_in_graph
-from pymc.pytensorf import convert_data
+from pymc.exceptions import ShapeError
+from pymc.pytensorf import convert_data, rvs_in_graph
 from pymc.vartypes import isgenerator
 
+if typing.TYPE_CHECKING:
+    from pymc.model.core import Model
+
 __all__ = [
     "Data",
     "Minibatch",
@@ -197,7 +199,7 @@ def determine_coords(
 
     if isinstance(value, np.ndarray) and dims is not None:
         if len(dims) != value.ndim:
-            raise pm.exceptions.ShapeError(
+            raise ShapeError(
                 "Invalid data shape. The rank of the dataset must match the length of `dims`.",
                 actual=value.shape,
                 expected=value.ndim,
@@ -222,6 +224,7 @@ def Data(
     dims: Sequence[str] | None = None,
     coords: dict[str, Sequence | np.ndarray] | None = None,
     infer_dims_and_coords=False,
+    model: Union["Model", None] = None,
     **kwargs,
 ) -> SharedVariable | TensorConstant:
     """Create a data container that registers a data variable with the model.
@@ -286,15 +289,18 @@ def Data(
     ...         model.set_data("data", data_vals)
     ...         idatas.append(pm.sample())
     """
+    from pymc.model.core import modelcontext
+
     if coords is None:
         coords = {}
 
     if isinstance(value, list):
         value = np.array(value)
 
     # Add data container to the named variables of the model.
-    model = pm.Model.get_context(error_if_none=False)
-    if model is None:
+    try:
+        model = modelcontext(model)
+    except TypeError:
         raise TypeError(
             "No model on context stack, which is needed to instantiate a data container. "
             "Add variable inside a 'with model:' block."
@@ -321,7 +327,7 @@ def Data(
     if isinstance(dims, str):
         dims = (dims,)
     if not (dims is None or len(dims) == x.ndim):
-        raise pm.exceptions.ShapeError(
+        raise ShapeError(
             "Length of `dims` must match the dimensions of the dataset.",
             actual=len(dims),
             expected=x.ndim,

pymc/dims/__init__.py

@@ -0,0 +1,67 @@
+#   Copyright 2025 - 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.
+
+
+def __init__():
+    """Make PyMC aware of the xtensor functionality.
+
+    This should be done eagerly once development matures.
+    """
+    import datetime
+    import warnings
+
+    from pytensor.compile import optdb
+
+    from pymc.initial_point import initial_point_rewrites_db
+    from pymc.logprob.abstract import MeasurableOp
+    from pymc.logprob.rewriting import logprob_rewrites_db
+
+    # Filter PyTensor xtensor warning, we emmit our own warning
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        import pytensor.xtensor
+
+        from pytensor.xtensor.vectorization import XRV
+
+    # Make PyMC aware of xtensor functionality
+    MeasurableOp.register(XRV)
+    lower_xtensor_query = optdb.query("+lower_xtensor")
+    logprob_rewrites_db.register("lower_xtensor", lower_xtensor_query, "basic", position=0.1)
+    initial_point_rewrites_db.register("lower_xtensor", lower_xtensor_query, "basic", position=0.1)
+
+    # TODO: Better model of probability of bugs
+    day_of_conception = datetime.date(2025, 6, 17)
+    day_of_last_bug = datetime.date(2025, 6, 30)
+    today = datetime.date.today()
+    days_with_bugs = (day_of_last_bug - day_of_conception).days
+    days_without_bugs = (today - day_of_last_bug).days
+    p = 1 - (days_without_bugs / (days_without_bugs + days_with_bugs + 10))
+    if p > 0.05:
+        warnings.warn(
+            f"The `pymc.dims` module is experimental and may contain critical bugs (p={p:.3f}).\n"
+            "Please report any issues you encounter at https://github.com/pymc-devs/pymc/issues.\n"
+            "Disclaimer: This an experimental API and may change at any time.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+
+__init__()
+del __init__
+
+from pytensor.xtensor import as_xtensor, concat
+
+from pymc.dims import math
+from pymc.dims.distributions import *
+from pymc.dims.model import Data, Deterministic, Potential, with_dims

pymc/dims/distributions/__init__.py

@@ -0,0 +1,15 @@
+#   Copyright 2025 - 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 pymc.dims.distributions.scalar import *
+from pymc.dims.distributions.vector import *