Enforce dims to be strings

Author: ricardoV94

pymc/backends/arviz.py

@@ -214,11 +214,9 @@ def __init__(
         }
 
         self.dims = {} if dims is None else dims
-        model_dims = {
-            var_name: [dim for dim in dims if dim is not None]
-            for var_name, dims in self.model.named_vars_to_dims.items()
-        }
+        model_dims = {k: list(v) for k, v in self.model.named_vars_to_dims.items()}
         self.dims = {**model_dims, **self.dims}
+
         if sample_dims is None:
             sample_dims = ["chain", "draw"]
         self.sample_dims = sample_dims

pymc/model.py

@@ -1317,6 +1317,8 @@ def register_rv(
         # the length of the corresponding RV dimension.
         if dims is not None:
             for d, dname in enumerate(dims):
+                if not isinstance(dname, str):
+                    raise TypeError(f"Dims must be string. Got {dname} of type {type(dname)}")
                 if dname not in self.dim_lengths:
                     self.add_coord(dname, values=None, length=rv_var.shape[d])
 

pymc/tests/test_model.py

@@ -871,6 +871,23 @@ def test_add_named_variable_checks_dim_name():
         pmodel.add_named_variable(rv2, dims=("nomnom", None))
 
 
+def test_dims_type_check():
+    with pm.Model(coords={"a": range(5)}) as m:
+        with pytest.raises(TypeError, match="Dims must be string"):
+            x = pm.Normal("x", shape=(10, 5), dims=(None, "a"))
+
+
+def test_none_coords_autonumbering():
+    with pm.Model() as m:
+        m.add_coord(name="a", values=None, length=3)
+        m.add_coord(name="b", values=range(5))
+        x = pm.Normal("x", dims=("a", "b"))
+        prior = pm.sample_prior_predictive(samples=2).prior
+    assert prior["x"].shape == (1, 2, 3, 5)
+    assert list(prior.coords["a"].values) == list(range(3))
+    assert list(prior.coords["b"].values) == list(range(5))
+
+
 def test_set_data_indirect_resize():
     with pm.Model() as pmodel:
         pmodel.add_coord("mdim", mutable=True, length=2)