Arviz don't fail hard on incompatible coordinate lengths

Author: ricardoV94

pymc/backends/arviz.py

@@ -49,10 +49,41 @@
 
 _log = logging.getLogger(__name__)
 
+
+RAISE_ON_INCOMPATIBLE_COORD_LENGTHS = False
+
+
 # random variable object ...
 Var = Any
 
 
+def dict_to_dataset_drop_incompatible_coords(vars_dict, *args, dims, coords, **kwargs):
+    safe_coords = coords
+
+    if not RAISE_ON_INCOMPATIBLE_COORD_LENGTHS:
+        coords_lengths = {k: len(v) for k, v in coords.items()}
+        for var_name, var in vars_dict.items():
+            # Iterate in reversed because of chain/draw batch dimensions
+            for dim, dim_length in zip(reversed(dims.get(var_name, ())), reversed(var.shape)):
+                coord_length = coords_lengths.get(dim, None)
+                if (coord_length is not None) and (coord_length != dim_length):
+                    warnings.warn(
+                        f"Incompatible coordinate length of {coord_length} for dimension '{dim}' of variable '{var_name}'.\n"
+                        "This usually happens when a sliced or concatenated variable is wrapped as a `pymc.dims.Deterministic`."
+                        "The originate coordinates for this dim will not be included in the returned dataset for any of the variables. "
+                        "Instead they will default to `np.arange(var_length)` and the shorter variables will be right-padded with nan.\n"
+                        "To make this warning into an error set `pymc.backends.arviz.RAISE_ON_INCOMPATIBLE_COORD_LENGTHS` to `True`",
+                        UserWarning,
+                    )
+                    if safe_coords is coords:
+                        safe_coords = coords.copy()
+                    safe_coords.pop(dim)
+                    coords_lengths.pop(dim)
+
+    # FIXME: Would be better to drop coordinates altogether, but arviz defaults to `np.arange(var_length)`
+    return dict_to_dataset(vars_dict, *args, dims=dims, coords=safe_coords, **kwargs)
+
+
 def find_observations(model: "Model") -> dict[str, Var]:
     """If there are observations available, return them as a dictionary."""
     observations = {}
@@ -365,7 +396,7 @@ def priors_to_xarray(self):
             priors_dict[group] = (
                 None
                 if var_names is None
-                else dict_to_dataset(
+                else dict_to_dataset_drop_incompatible_coords(
                     {k: np.expand_dims(self.prior[k], 0) for k in var_names},
                     library=pymc,
                     coords=self.coords,

tests/backends/test_arviz.py

@@ -11,6 +11,7 @@
 #   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.
+import re
 import warnings
 
 import numpy as np
@@ -837,3 +838,27 @@ def test_dataset_to_point_list_str_key(self):
         ds[3] = xarray.DataArray([1, 2, 3])
         with pytest.raises(ValueError, match="must be str"):
             dataset_to_point_list(ds, sample_dims=["chain", "draw"])
+
+
+def test_incompatible_coordinate_lengths():
+    with pm.Model(coords={"a": [-1, -2, -3]}) as m:
+        x = pm.Normal("x", dims="a")
+        y = pm.Deterministic("y", x[1:], dims=("a",))
+
+        with pytest.warns(
+            UserWarning,
+            match=re.escape(
+                "Incompatible coordinate length of 3 for dimension 'a' of variable 'y'"
+            ),
+        ):
+            prior = pm.sample_prior_predictive(draws=1).prior.squeeze(("chain", "draw"))
+        assert prior.x.dims == prior.y.dims == ("a",)
+        assert prior.x.shape == prior.y.shape == (3,)
+        assert np.isnan(prior.y.values[-1])
+        assert list(prior.coords["a"]) == [0, 1, 2]
+
+        pm.backends.arviz.RAISE_ON_INCOMPATIBLE_COORD_LENGTHS = True
+        with pytest.raises(ValueError):
+            pm.sample_prior_predictive(draws=1)
+
+        pm.backends.arviz.RAISE_ON_INCOMPATIBLE_COORD_LENGTHS = False