Automatically resize independently mutable dims via set_data

Author: michaelosthege

pymc/model.py

@@ -1199,28 +1199,29 @@ def set_data(
                     # definitely lead to shape problems.
                     raise ShapeError(
                         f"Resizing dimension '{dname}' is impossible, because "
-                        "a 'TensorConstant' stores its length. To be able "
+                        "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.
+                elif length_tensor.owner is not None:
+                    # The dimension was created from a model variable.
                     length_belongs_to = length_tensor.owner.inputs[0].owner.inputs[0]
-                    if not isinstance(length_belongs_to, SharedVariable):
+                    if length_belongs_to is shared_object:
+                        # No surprise it's changing.
+                        pass
+                    elif isinstance(length_belongs_to, SharedVariable):
+                        # The dimension is mutable through a SharedVariable other than the one being modified.
+                        # But the other variable was not yet re-sized! Warn the user to do that!
+                        warnings.warn(
+                            f"You are resizing a variable with dimension '{dname}' which was initialized "
+                            f"as a mutable dimension by another variable ('{length_belongs_to}')."
+                            " Remember to update that variable with the correct shape to avoid shape issues.",
+                            ShapeWarning,
+                            stacklevel=2,
+                        )
+                    else:
+                        # The dimension is immutable.
                         raise ShapeError(
                             f"Resizing dimension '{dname}' with values of length {new_length} would lead to incompatibilities, "
                             f"because the dimension was initialized from '{length_belongs_to}' which is not a shared variable. "
@@ -1230,8 +1231,9 @@ def set_data(
                             expected=old_length,
                         )
                 if isinstance(length_tensor, ScalarSharedVariable):
-                    # Updating the shared variable resizes dependent nodes that use this dimension for their `size`.
-                    length_tensor.set_value(new_length)
+                    # The dimension is mutable, but was defined without being linked
+                    # to a shared variable. This is allowed, but a little less robust.
+                    self.set_dim(dname, new_length, coord_values=new_coords)
 
             if new_coords is not None:
                 # Update the registered coord values (also if they were None)

pymc/tests/test_model.py

@@ -12,6 +12,7 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 import unittest
+import warnings
 
 from functools import reduce
 
@@ -704,25 +705,45 @@ def test_nested_model_coords():
     assert set(m2.RV_dims) < set(m1.RV_dims)
 
 
-def test_shapeerror_from_resize_immutable_dims():
+def test_shapeerror_from_set_data_dimensionality():
+    with pm.Model() as pmodel:
+        pm.MutableData("m", np.ones((3,)), dims="one")
+        with pytest.raises(ValueError, match="must have 1 dimensions"):
+            pmodel.set_data("m", np.ones((3, 4)))
+
+
+def test_shapeerror_from_resize_immutable_dim_from_RV():
     """
     Trying to resize an immutable dimension should raise a ShapeError.
     Even if the variable being updated is a SharedVariable and has other
     dimensions that are mutable.
     """
     with pm.Model() as pmodel:
-        a = pm.Normal("a", mu=[1, 2, 3], dims="fixed")
+        pm.Normal("a", mu=[1, 2, 3], dims="fixed")
+        assert isinstance(pmodel.dim_lengths["fixed"], TensorVariable)
 
-        m = pm.MutableData("m", [[1, 2, 3]], dims=("one", "fixed"))
+        pm.MutableData("m", [[1, 2, 3]], dims=("one", "fixed"))
 
         # This is fine because the "fixed" dim is not resized
-        pm.set_data({"m": [[1, 2, 3], [3, 4, 5]]})
+        pmodel.set_data("m", [[1, 2, 3], [3, 4, 5]])
 
     with pytest.raises(ShapeError, match="was initialized from 'a'"):
-        # Can't work because the "fixed" dimension is linked to a constant shape:
+        # Can't work because the "fixed" dimension is linked to a
+        # TensorVariable with constant shape.
         # Note that the new data tries to change both dimensions
-        with pmodel:
-            pm.set_data({"m": [[1, 2], [3, 4]]})
+        pmodel.set_data("m", [[1, 2], [3, 4]])
+
+
+def test_shapeerror_from_resize_immutable_dim_from_coords():
+    with pm.Model(coords={"immutable": [1, 2]}) as pmodel:
+        assert isinstance(pmodel.dim_lengths["immutable"], TensorConstant)
+        pm.MutableData("m", [1, 2], dims="immutable")
+        # Data can be changed
+        pmodel.set_data("m", [3, 4])
+
+    with pytest.raises(ShapeError, match="`TensorConstant` stores its length"):
+        # But the length is linked to a TensorConstant
+        pmodel.set_data("m", [1, 2, 3], coords=dict(immutable=[1, 2, 3]))
 
 
 def test_valueerror_from_resize_without_coords_update():
@@ -798,22 +819,65 @@ def test_set_dim_with_coords():
     assert pmodel.coords["mdim"] == ("A", "B", "C")
 
 
-def test_set_data_warns_resize_mutable_dim():
+def test_set_data_indirect_resize():
     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])
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        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"):
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
         pmodel.set_data("mdata", [1, 2, 3, 4])
-    pass
+
+
+def test_set_data_warns_on_resize_of_dims_defined_by_other_mutabledata():
+    with pm.Model() as pmodel:
+        pm.MutableData("m1", [1, 2], dims="mutable")
+        pm.MutableData("m2", [3, 4], dims="mutable")
+
+        # Resizing the non-defining variable first gives a warning
+        with pytest.warns(ShapeWarning, match="by another variable"):
+            pmodel.set_data("m2", [4, 5, 6])
+            pmodel.set_data("m1", [1, 2, 3])
+
+        # Resizing the definint variable first is silent
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            pmodel.set_data("m1", [1, 2])
+            pmodel.set_data("m2", [3, 4])
+
+
+def test_set_data_indirect_resize_with_coords():
+    with pm.Model() as pmodel:
+        pmodel.add_coord("mdim", ["A", "B"], mutable=True, length=2)
+        pm.MutableData("mdata", [1, 2], dims="mdim")
+
+    assert pmodel.coords["mdim"] == ("A", "B")
+
+    # First resize the dimension.
+    pmodel.set_dim("mdim", 3, ["A", "B", "C"])
+    assert pmodel.coords["mdim"] == ("A", "B", "C")
+    # Then change the data.
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        pmodel.set_data("mdata", [1, 2, 3])
+
+    # Now the other way around.
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        pmodel.set_data("mdata", [1, 2, 3, 4], coords=dict(mdim=["A", "B", "C", "D"]))
+    assert pmodel.coords["mdim"] == ("A", "B", "C", "D")
+
+    # This time with incorrectly sized coord values
+    with pytest.raises(ShapeError, match="new coordinate values"):
+        pmodel.set_data("mdata", [1, 2], coords=dict(mdim=[1, 2, 3]))
 
 
 @pytest.mark.parametrize("jacobian", [True, False])