Warn about resizing MutableData dims that are not symbolically linked

Closes #5812

Author: michaelosthege

pymc/model.py

@@ -60,7 +60,7 @@
 from pymc.distributions import joint_logpt
 from pymc.distributions.logprob import _get_scaling
 from pymc.distributions.transforms import _default_transform
-from pymc.exceptions import ImputationWarning, SamplingError, ShapeError
+from pymc.exceptions import ImputationWarning, SamplingError, ShapeError, ShapeWarning
 from pymc.initial_point import make_initial_point_fn
 from pymc.math import flatten_list
 from pymc.util import (
@@ -1193,14 +1193,31 @@ def set_data(
                             f"{new_length}, so new coord values for the {dname} dimension are required."
                         )
                 if isinstance(length_tensor, TensorConstant):
+                    # The dimension was fixed in length.
+                    # Resizing a data variable in this dimension would
+                    # definitely lead to shape problems.
                     raise ShapeError(
                         f"Resizing dimension '{dname}' is impossible, because "
                         "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."
                     )
+                elif isinstance(length_tensor, ScalarSharedVariable):
+                    # The dimension is mutable, but was defined without being linked
+                    # to a shared variable. This is allowed, but slightly dangerous.
+                    warnings.warn(
+                        f"You are resizing a variable with dimension '{dname}' which was initialized"
+                        " as a mutable dimension and is not linked to the `MutableData` variable."
+                        " Remember to update the dimension length by calling "
+                        f"`Model.set_dim({dname}, new_length={new_length})` manually,"
+                        " preferably _before_ updating `MutableData` variables that use this dimension.",
+                        ShapeWarning,
+                        stacklevel=2,
+                    )
                 else:
+                    # The dimension was created from another model variable.
+                    # If that was a non-mutable variable, there will definitely be shape problems.
                     length_belongs_to = length_tensor.owner.inputs[0].owner.inputs[0]
                     if not isinstance(length_belongs_to, SharedVariable):
                         raise ShapeError(

pymc/tests/test_model.py

@@ -37,7 +37,7 @@
 from pymc import Deterministic, Potential
 from pymc.blocking import DictToArrayBijection, RaveledVars
 from pymc.distributions import Normal, transforms
-from pymc.exceptions import ShapeError
+from pymc.exceptions import ShapeError, ShapeWarning
 from pymc.model import Point, ValueGradFunction
 from pymc.tests.helpers import SeededTest
 
@@ -798,6 +798,24 @@ def test_set_dim_with_coords():
     assert pmodel.coords["mdim"] == ("A", "B", "C")
 
 
+def test_set_data_warns_resize_mutable_dim():
+    with pm.Model() as pmodel:
+        pmodel.add_coord("mdim", mutable=True, length=2)
+        pm.MutableData("mdata", [1, 2], dims="mdim")
+
+    # First resize the dimension.
+    pmodel.dim_lengths["mdim"].set_value(3)
+    # Then change the data.
+    pmodel.set_data("mdata", [1, 2, 3])
+
+    # Now the other way around.
+    # Because the dimension doesn't depend on the data variable,
+    # a warning shoudl be emitted.
+    with pytest.warns(ShapeWarning, match="update the dimension length"):
+        pmodel.set_data("mdata", [1, 2, 3, 4])
+    pass
+
+
 @pytest.mark.parametrize("jacobian", [True, False])
 def test_model_logp(jacobian):
     with pm.Model() as m: