Use dict_to_dataset_drop_incompatible_coords everywhere

Author: ricardoV94

pymc/backends/arviz.py

@@ -303,14 +303,14 @@ def posterior_to_xarray(self):
                     self.posterior_trace.get_values(var_name, combine=False, squeeze=False)
                 )
         return (
-            dict_to_dataset(
+            dict_to_dataset_drop_incompatible_coords(
                 data,
                 library=pymc,
                 coords=self.coords,
                 dims=self.dims,
                 attrs=self.attrs,
             ),
-            dict_to_dataset(
+            dict_to_dataset_drop_incompatible_coords(
                 data_warmup,
                 library=pymc,
                 coords=self.coords,
@@ -345,14 +345,14 @@ def sample_stats_to_xarray(self):
                 )
 
         return (
-            dict_to_dataset(
+            dict_to_dataset_drop_incompatible_coords(
                 data,
                 library=pymc,
                 dims=None,
                 coords=self.coords,
                 attrs=self.attrs,
             ),
-            dict_to_dataset(
+            dict_to_dataset_drop_incompatible_coords(
                 data_warmup,
                 library=pymc,
                 dims=None,
@@ -366,7 +366,7 @@ def posterior_predictive_to_xarray(self):
         """Convert posterior_predictive samples to xarray."""
         data = self.posterior_predictive
         dims = {var_name: self.sample_dims + self.dims.get(var_name, []) for var_name in data}
-        return dict_to_dataset(
+        return dict_to_dataset_drop_incompatible_coords(
             data, library=pymc, coords=self.coords, dims=dims, default_dims=self.sample_dims
         )
 
@@ -375,7 +375,7 @@ def predictions_to_xarray(self):
         """Convert predictions (out of sample predictions) to xarray."""
         data = self.predictions
         dims = {var_name: self.sample_dims + self.dims.get(var_name, []) for var_name in data}
-        return dict_to_dataset(
+        return dict_to_dataset_drop_incompatible_coords(
             data, library=pymc, coords=self.coords, dims=dims, default_dims=self.sample_dims
         )
 
@@ -412,7 +412,7 @@ def observed_data_to_xarray(self):
         """Convert observed data to xarray."""
         if self.predictions:
             return None
-        return dict_to_dataset(
+        return dict_to_dataset_drop_incompatible_coords(
             self.observations,
             library=pymc,
             coords=self.coords,
@@ -427,7 +427,7 @@ def constant_data_to_xarray(self):
         if not constant_data:
             return None
 
-        xarray_dataset = dict_to_dataset(
+        xarray_dataset = dict_to_dataset_drop_incompatible_coords(
             constant_data,
             library=pymc,
             coords=self.coords,
@@ -705,7 +705,7 @@ def apply_function_over_dataset(
             )
         )
 
-    return dict_to_dataset(
+    return dict_to_dataset_drop_incompatible_coords(
         out_trace,
         library=pymc,
         dims=dims,

pymc/sampling/mcmc.py

@@ -33,7 +33,7 @@
 import numpy as np
 import pytensor.gradient as tg
 
-from arviz import InferenceData, dict_to_dataset
+from arviz import InferenceData
 from arviz.data.base import make_attrs
 from pytensor.graph.basic import Variable
 from rich.theme import Theme
@@ -45,6 +45,7 @@
 from pymc.backends import RunType, TraceOrBackend, init_traces
 from pymc.backends.arviz import (
     coords_and_dims_for_inferencedata,
+    dict_to_dataset_drop_incompatible_coords,
     find_constants,
     find_observations,
 )
@@ -355,14 +356,14 @@ def _sample_external_nuts(
         # Temporary work-around. Revert once https://github.com/pymc-devs/nutpie/issues/74 is fixed
         # gather observed and constant data as nutpie.sample() has no access to the PyMC model
         coords, dims = coords_and_dims_for_inferencedata(model)
-        constant_data = dict_to_dataset(
+        constant_data = dict_to_dataset_drop_incompatible_coords(
             find_constants(model),
             library=pm,
             coords=coords,
             dims=dims,
             default_dims=[],
         )
-        observed_data = dict_to_dataset(
+        observed_data = dict_to_dataset_drop_incompatible_coords(
             find_observations(model),
             library=pm,
             coords=coords,

pymc/smc/sampling.py

@@ -33,7 +33,7 @@
 
 import pymc
 
-from pymc.backends.arviz import dict_to_dataset, to_inference_data
+from pymc.backends.arviz import dict_to_dataset_drop_incompatible_coords, to_inference_data
 from pymc.backends.base import MultiTrace
 from pymc.distributions.custom import CustomDistRV, CustomSymbolicDistRV
 from pymc.distributions.distribution import _support_point
@@ -264,7 +264,7 @@ def _save_sample_stats(
             else:
                 sample_stats_dict[stat] = np.array(value)
 
-        sample_stats = dict_to_dataset(
+        sample_stats = dict_to_dataset_drop_incompatible_coords(
             sample_stats_dict,
             attrs=sample_settings_dict,
             library=pymc,