From c535718e65aa579fcf525fac2bcad44219e9c046 Mon Sep 17 00:00:00 2001
From: Cove Geary <covertgeary@gmail.com>
Date: Sun, 12 Mar 2023 17:20:46 -0500
Subject: [PATCH] Fix bug introduced in arviz.py and add regression test for
 #6496

---
 pymc/backends/arviz.py       | 16 ++++++++++------
 pymc/model.py                |  5 ++---
 tests/backends/test_arviz.py | 33 ++++++++++++++++++++++++---------
 3 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/pymc/backends/arviz.py b/pymc/backends/arviz.py
index fcd63e386f..2602c6b920 100644
--- a/pymc/backends/arviz.py
+++ b/pymc/backends/arviz.py
@@ -216,15 +216,19 @@ def __init__(
                 " one of trace, prior, posterior_predictive or predictions."
             )
 
-        coords_typed = self.model.coords_typed
-        if coords:
-            coords_typed.update(_as_coord_vals(coords))
-        coords_typed = {
+        given_coords = coords if coords is not None else {}
+        given_coords_typed = {
+            cname: _as_coord_vals(cvals)
+            for cname, cvals in given_coords.items()
+            if cvals is not None
+        }
+        model_coords_typed = {
             cname: cvals_typed
-            for cname, cvals_typed in coords_typed.items()
+            for cname, cvals_typed in self.model.coords_typed.items()
             if cvals_typed is not None
         }
-        self.coords = coords_typed
+        # Coords from argument should have precedence
+        self.coords = {**model_coords_typed, **given_coords_typed}
 
         self.dims = {} if dims is None else dims
         model_dims = {k: list(v) for k, v in self.model.named_vars_to_dims.items()}
diff --git a/pymc/model.py b/pymc/model.py
index ea65a9fc6b..bb05c197d6 100644
--- a/pymc/model.py
+++ b/pymc/model.py
@@ -1061,9 +1061,8 @@ def add_coord(
         if values is not None:
             # Conversion to numpy array to ensure coord vals are 1-dim
             values = _as_coord_vals(values)
-        if name in self.coords_typed:
-            if not np.array_equal(_as_coord_vals(values), self.coords_typed[name]):
-                raise ValueError(f"Duplicate and incompatible coordinate: {name}.")
+        if name in self.coords_typed and not np.array_equal(values, self.coords_typed[name]):
+            raise ValueError(f"Duplicate and incompatible coordinate: {name}.")
         if length is not None and not isinstance(length, (int, Variable)):
             raise ValueError(
                 f"The `length` passed for the '{name}' coord must be an int, PyTensor Variable or None."
diff --git a/tests/backends/test_arviz.py b/tests/backends/test_arviz.py
index 73c3e9b5cf..f6074abf6e 100644
--- a/tests/backends/test_arviz.py
+++ b/tests/backends/test_arviz.py
@@ -630,15 +630,30 @@ def test_issue_5043_autoconvert_coord_values(self):
             # We're not automatically converting things other than tuple,
             # so advanced use cases remain supported at the InferenceData level.
             # They just can't be used in the model construction already.
-            converter = InferenceDataConverter(
-                trace=mtrace,
-                coords={
-                    "city": pd.MultiIndex.from_tuples(
-                        [("Bonn", 53111), ("Berlin", 10178)], names=["name", "zipcode"]
-                    )
-                },
-            )
-            assert isinstance(converter.coords["city"], pd.MultiIndex)
+            # TODO: we now ensure that everything passed to arviz is converted to
+            # ndarray, which makes the following test fail.
+            # converter = InferenceDataConverter(
+            #     trace=mtrace,
+            #     coords={
+            #         "city": pd.MultiIndex.from_tuples(
+            #             [("Bonn", 53111), ("Berlin", 10178)], names=["name", "zipcode"]
+            #         )
+            #     },
+            # )
+            # assert isinstance(converter.coords["city"], pd.MultiIndex)
+
+    def test_nested_coords_issue_6496(self):
+        """Regression test to ensure we don't bug out if coordinate values
+        appear "nested" to numpy.
+        """
+        model = pm.Model(coords={"cname": [("a", 1), ("a", 2), ("b", 1)]})
+        idata = to_inference_data(
+            prior={"x": np.zeros((100, 3))}, dims={"x": ["cname"]}, model=model
+        )
+        idata_coord = idata.prior.coords["cname"]
+        assert len(idata_coord) == 3
+        assert idata_coord.dtype == np.dtype("O")
+        assert np.array_equal(idata_coord.data, model.coords_typed["cname"])
 
     def test_variable_dimension_name_collision(self):
         with pytest.raises(ValueError, match="same name as its dimension"):