fix(pymc): Allow dims with only length specified

Author: aseyboldt

python/nutpie/compile_pymc.py

@@ -207,7 +207,12 @@ def compile_pymc_model(model: "pm.Model", **kwargs) -> CompiledPyMCModel:
     )
     expand_numba = numba.cfunc(c_sig_expand, **kwargs)(expand_numba_raw)
 
-    coords = {name: pd.Index(vals) for name, vals in model.coords.items()}
+    coords = {}
+    for name, vals in model.coords.items():
+        if vals is None:
+            vals = pd.RangeIndex(int(model.dim_lengths[name].eval()))
+        coords[name] = pd.Index(vals)
+
     if "unconstrained_parameter" in coords:
         raise ValueError("Model contains invalid name 'unconstrained_parameter'.")
 

tests/test_pymc.py

@@ -15,6 +15,16 @@ def test_pymc_model():
     trace.posterior.a
 
 
+def test_pymc_model():
+    with pm.Model() as model:
+        model.add_coord("foo", length=5)
+        pm.Normal("a", dims="foo")
+
+    compiled = nutpie.compile_pymc_model(model)
+    trace = nutpie.sample(compiled, chains=1)
+    trace.posterior.a
+
+
 def test_trafo():
     with pm.Model() as model:
         pm.Uniform("a")