Fix xarray DataArray coord unwrapping in Model.add_coord (#8391)

When xarray Dataset coords are passed directly to pm.Model(coords=...),
the DataArray values were stored as tuples of 0-d DataArrays instead of
plain values. This broke np.array_equal comparisons and InferenceData
export.

Fix by extracting .values from any DataArray before tuple() conversion.

Author: williambdean

pymc/model/core.py

@@ -29,6 +29,7 @@
 import numpy as np
 import pytensor
 import pytensor.tensor as pt
+import xarray as xr
 
 from pytensor.compile import DeepCopyOp, Function, ProfileStats, get_mode, view_op
 from pytensor.compile.io import In, Out
@@ -1856,6 +1857,12 @@ def add_coord(
                 f"Either `values` or `length` must be specified for the '{name}' dimension."
             )
         if values is not None:
+            # xarray DataArrays passed as coord values must be unwrapped.
+            # tuple(DataArray) iterates yielding 0-d DataArrays instead of plain values,
+            # which breaks np.array_equal comparisons and DataTree export.
+            if isinstance(values, xr.DataArray):
+                values = values.values
+
             # Conversion to a tuple ensures that the coordinate values are immutable.
             # Also unlike numpy arrays the's tuple.index(...) which is handy to work with.
             values = tuple(values)

tests/model/test_core.py

@@ -31,6 +31,7 @@
 import scipy
 import scipy.sparse as sps
 import scipy.stats as st
+import xarray as xr
 
 from pytensor.compile.mode import get_default_mode
 from pytensor.compile.sharedvalue import SharedVariable
@@ -873,6 +874,27 @@ def test_multiple_add_coords_with_same_name():
     assert len(variables) == 1 and variables[0] is m.dim_lengths["dim1"]
 
 
+@pytest.mark.parametrize(
+    "coords_dict",
+    [
+        pytest.param({"city": ["nyc", "la", "chi"]}, id="string"),
+        pytest.param({"year": [2020, 2021, 2022]}, id="int"),
+        pytest.param(
+            {"time": np.array(["2020-01-01", "2020-01-02"], dtype="datetime64[D]")},
+            id="datetime64",
+        ),
+    ],
+)
+def test_xarray_coord_values_unwrapped(coords_dict):
+    """xarray DataArray coord values are unwrapped to plain tuples of values."""
+    ds = xr.Dataset(coords=coords_dict)
+    with pm.Model(coords=ds.coords) as m:
+        key = next(iter(coords_dict))
+        coord = m.coords[key]
+        assert isinstance(coord, tuple)
+        assert not isinstance(coord[0], xr.DataArray)
+
+
 class TestSetUpdateCoords:
     def test_shapeerror_from_set_data_dimensionality(self):
         with pm.Model() as pmodel: