improve coords matching (#623)

pymc-labs · Apr 11, 2024 · a151330 · a151330
1 parent ffef5a6
commit a151330
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/tests/clv/models/test_shifted_beta_geo.py b/tests/clv/models/test_shifted_beta_geo.py
@@ -183,7 +183,7 @@ def test_model_convergence(self, data, model_config):
     def test_distribution_customer_churn_time(self):
         dataset = pd.DataFrame(
             {
-                "customer_id": [1, 2, 3],
+                "customer_id": [0, 1, 2],
                 "t_churn": [10, 10, 10],
                 "T": 10,
             }
@@ -195,11 +195,13 @@ def test_distribution_customer_churn_time(self):
         model.fit(fit_method="map")
         customer_thetas = np.array([0.1, 0.5, 0.9])
         model.idata = az.from_dict(
-            {
+            posterior={
                 "alpha": np.ones((2, 500)),  # Two chains, 500 draws each
                 "beta": np.ones((2, 500)),
                 "theta": np.full((2, 500, 3), customer_thetas),
-            }
+            },
+            coords={"customer_id": [0, 1, 2]},
+            dims={"theta": ["customer_id"]},
         )
 
         res = model.distribution_customer_churn_time(