Require pm.MutableData or add_coord(mutable=True) for mutable dims

This increases the safety of working with resizable dimensions.
Dims created via `pm.Model(coords=...)` coordinates are now immutable.
Only dimension lengths that are created anew from `pm.MutableData`,
`pm.Data(mutable=True)`, or the underlying `add_coord(mutable=True)`
will become shared variables.

Co-authored-by: Michael Osthege <[email protected]>

Author: Luke LB

pymc/data.py

@@ -704,6 +704,7 @@ def Data(
                     # Note: Coordinate values can't be taken from
                     # the value, because it could be N-dimensional.
                     values=coords.get(dname, None),
+                    mutable=mutable,
                     length=xshape[d],
                 )
 

pymc/model.py

@@ -1071,6 +1071,7 @@ def add_coord(
         self,
         name: str,
         values: Optional[Sequence] = None,
+        mutable: bool = False,
         *,
         length: Optional[Union[int, Variable]] = None,
     ):
@@ -1084,9 +1085,12 @@ def add_coord(
         values : optional, array-like
             Coordinate values or ``None`` (for auto-numbering).
             If ``None`` is passed, a ``length`` must be specified.
+        mutable : bool
+            Whether the created dimension should be resizable.
+            Default is False.
         length : optional, scalar
             A scalar of the dimensions length.
-            Defaults to ``aesara.shared(len(values))``.
+            Defaults to ``aesara.tensor.constant(len(values))``.
         """
         if name in {"draw", "chain", "__sample__"}:
             raise ValueError(
@@ -1111,8 +1115,11 @@ def add_coord(
             if not np.array_equal(values, self.coords[name]):
                 raise ValueError(f"Duplicate and incompatible coordinate: {name}.")
         else:
+            if mutable:
+                self._dim_lengths[name] = length or aesara.shared(len(values))
+            else:
+                self._dim_lengths[name] = length or aesara.tensor.constant(len(values))
             self._coords[name] = values
-            self._dim_lengths[name] = length or aesara.shared(len(values))
 
     def add_coords(
         self,
@@ -1192,9 +1199,10 @@ def set_data(
                 if isinstance(length_tensor, TensorConstant):
                     raise ShapeError(
                         f"Resizing dimension '{dname}' is impossible, because "
-                        f"a 'TensorConstant' stores its length. To be able "
-                        f"to change the dimension length, 'fixed' in "
-                        f"'model.add_coord' must be set to `False`."
+                        "a 'TensorConstant' stores its length. To be able "
+                        "to change the dimension length, pass `mutable=True` when "
+                        "registering the dimension via `model.add_coord`, "
+                        "or define it via a `pm.MutableData` variable."
                     )
                 else:
                     length_belongs_to = length_tensor.owner.inputs[0].owner.inputs[0]

pymc/tests/test_data_container.py

@@ -19,7 +19,7 @@
 
 from aesara import shared
 from aesara.compile.sharedvalue import SharedVariable
-from aesara.tensor.sharedvar import ScalarSharedVariable
+from aesara.tensor import TensorConstant
 from aesara.tensor.var import TensorVariable
 
 import pymc as pm
@@ -315,21 +315,22 @@ def test_explicit_coords(self):
         }
         # pass coordinates explicitly, use numpy array in Data container
         with pm.Model(coords=coords) as pmodel:
+            # Dims created from coords are constant by default
+            assert isinstance(pmodel.dim_lengths["rows"], TensorConstant)
+            assert isinstance(pmodel.dim_lengths["columns"], TensorConstant)
             pm.MutableData("observations", data, dims=("rows", "columns"))
-            # new data with same shape
+            # new data with same (!) shape
             pm.set_data({"observations": data + 1})
-            # new data with same shape and coords
+            # new data with same (!) shape and coords
             pm.set_data({"observations": data}, coords=coords)
         assert "rows" in pmodel.coords
         assert pmodel.coords["rows"] == ("R1", "R2", "R3", "R4", "R5")
         assert "rows" in pmodel.dim_lengths
-        assert isinstance(pmodel.dim_lengths["rows"], ScalarSharedVariable)
         assert pmodel.dim_lengths["rows"].eval() == 5
         assert "columns" in pmodel.coords
         assert pmodel.coords["columns"] == ("C1", "C2", "C3", "C4", "C5", "C6", "C7")
         assert pmodel.RV_dims == {"observations": ("rows", "columns")}
         assert "columns" in pmodel.dim_lengths
-        assert isinstance(pmodel.dim_lengths["columns"], ScalarSharedVariable)
         assert pmodel.dim_lengths["columns"].eval() == 7
 
     def test_set_coords_through_pmdata(self):
@@ -354,6 +355,7 @@ def test_symbolic_coords(self):
         Their lengths are then automatically linked to the corresponding Tensor dimension.
         """
         with pm.Model() as pmodel:
+            # Dims created from MutableData are TensorVariables linked to the SharedVariable.shape
             intensity = pm.MutableData("intensity", np.ones((2, 3)), dims=("row", "column"))
             assert "row" in pmodel.dim_lengths
             assert "column" in pmodel.dim_lengths

pymc/tests/test_model.py

@@ -27,7 +27,9 @@
 import scipy.sparse as sps
 import scipy.stats as st
 
+from aesara.tensor import TensorVariable
 from aesara.tensor.random.op import RandomVariable
+from aesara.tensor.sharedvar import ScalarSharedVariable
 from aesara.tensor.var import TensorConstant
 
 import pymc as pm
@@ -729,13 +731,40 @@ def test_valueerror_from_resize_without_coords_update():
     without passing new coords raises a ValueError.
     """
     with pm.Model() as pmodel:
-        pmodel.add_coord("shared", [1, 2, 3])
+        pmodel.add_coord("shared", [1, 2, 3], mutable=True)
         pm.MutableData("m", [1, 2, 3], dims=("shared"))
         with pytest.raises(ValueError, match="'m' variable already had 3"):
             # tries to resize m but without passing coords so raise ValueError
             pm.set_data({"m": [1, 2, 3, 4]})
 
 
+def test_coords_and_constantdata_create_immutable_dims():
+    """
+    When created from `pm.Model(coords=...)` or `pm.ConstantData`
+    a dimension should be resizable.
+    """
+    with pm.Model(coords={"group": ["A", "B"]}) as m:
+        x = pm.ConstantData("x", [0], dims="feature")
+        y = pm.Normal("y", x, 1, dims=("group", "feature"))
+    assert isinstance(m._dim_lengths["feature"], TensorConstant)
+    assert isinstance(m._dim_lengths["group"], TensorConstant)
+    assert x.eval().shape == (1,)
+    assert y.eval().shape == (2, 1)
+
+
+def test_add_coord_mutable_kwarg():
+    """
+    Checks resulting tensor type depending on mutable kwarg in add_coord.
+    """
+    with pm.Model() as m:
+        m.add_coord("fixed", values=[1], mutable=False)
+        m.add_coord("mutable1", values=[1, 2], mutable=True)
+        assert isinstance(m._dim_lengths["fixed"], TensorConstant)
+        assert isinstance(m._dim_lengths["mutable1"], ScalarSharedVariable)
+        pm.MutableData("mdata", np.ones((1, 2, 3)), dims=("fixed", "mutable1", "mutable2"))
+        assert isinstance(m._dim_lengths["mutable2"], TensorVariable)
+
+
 @pytest.mark.parametrize("jacobian", [True, False])
 def test_model_logp(jacobian):
     with pm.Model() as m: