Arviz don't fail hard on incompatible coordinates

Author: ricardoV94

pymc/backends/arviz.py

@@ -49,10 +49,40 @@
 
 _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} found for dimension {dim} of variable {var_name}.\n"
+                        "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)`.\n"
+                        "To make this warning into an errror 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 +395,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,