Add simple test case for coords behavior (#3977)

* Add simple test case for coords behavior * Install arviz directly from github as temporary fix * Revert attempt to point to latest arviz github * Bump arviz requirement Co-authored-by: Michael Osthege <[email protected]>
pymc-devs · Jun 24, 2020 · 9c7ed77 · 9c7ed77
1 parent 69b7975
commit 9c7ed77
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/pymc3/tests/test_coords.py b/pymc3/tests/test_coords.py
@@ -0,0 +1,18 @@
+import pymc3 as pm
+import numpy as np
+
+
+def test_coords():
+    chains = 2
+    n_features = 3
+    n_samples = 10
+
+    coords = {"features": np.arange(n_features)}
+
+    with pm.Model(coords=coords):
+        a = pm.Uniform("a", -100, 100, dims="features")
+        b = pm.Uniform("b", -100, 100, dims="features")
+        tr = pm.sample(n_samples, chains=chains, return_inferencedata=True)
+
+    assert "features" in tr.posterior.a.coords.dims
+    assert "features" in tr.posterior.b.coords.dims
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,4 @@
-arviz>=0.8.3
+arviz>=0.9.0
 theano>=1.0.4
 numpy>=1.13.0
 scipy>=0.18.1