Repository: pymc-devs/pymc

Issue:
Fix dim lengths created from `coords` or `ConstantData`
In [this comment](https://github.com/pymc-devs/pymc/issues/5613#issuecomment-1107752199) we proposed to make dimensions created from `coords` non-resizeable, because it would prevent users from running into issues like #5181.

## What needs to be done
In https://github.com/pymc-devs/pymc/blob/dd23c9008302fc184298076669d0a3a332aa59f0/pymc/model.py#L1070-L1076 the `set_coord` method could take another kwarg `resizable: bool=False` or `fixed: bool=True`.

Then we can actively turn off re-sizing of dimensions that are created from `ConstantData` or `coords`.

Discussion:
canyon289: > it's worth seeing if they fail on here and then look at fixes, which will be faster than me trying to figure things out on my own

Great strategy. Thanks for taking such a principled approach to development

It seems that there type being changed somewhere. Might be as simple as updating this test.
https://github.com/pymc-devs/pymc/runs/6435716258?check_suite_focus=true#step:7:80
codecov[bot]: # [Codecov](https://codecov.io/gh/pymc-devs/pymc/pull/5763?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
> Merging [#5763](https://codecov.io/gh/pymc-devs/pymc/pull/5763?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (c45a7e9) into [main](https://codecov.io/gh/pymc-devs/pymc/commit/ab593b114ad4491da5e9082abf3095fd9c03f169?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (ab593b1) will **decrease** coverage by `13.76%`.
> The diff coverage is `66.66%`.

[![Impacted file tree graph](https://codecov.io/gh/pymc-devs/pymc/pull/5763/graphs/tree.svg?width=650&height=150&src=pr&token=JFuXtOJ4Cb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://codecov.io/gh/pymc-devs/pymc/pull/5763?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

```diff
@@             Coverage Diff             @@
##             main    #5763       +/-   ##
===========================================
- Coverage   89.27%   75.50%   -13.77%     
===========================================
  Files          74       74               
  Lines       13803    13805        +2     
===========================================
- Hits        12322    10424     -1898     
- Misses       1481     3381     +1900     
```


| [Impacted Files](https://codecov.io/gh/pymc-devs/pymc/pull/5763?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pymc/model.py](https://codecov.io/gh/pymc-devs/pymc/pull/5763/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9tb2RlbC5weQ==) | `72.77% <66.66%> (-13.78%)` | :arrow_down: |
| [pymc/model\_graph.py](https://codecov.io/gh/pymc-devs/pymc/pull/5763/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9tb2RlbF9ncmFwaC5weQ==) | `15.86% <0.00%> (-69.66%)` | :arrow_down: |
| [pymc/func\_utils.py](https://codecov.io/gh/pymc-devs/pymc/pull/5763/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9mdW5jX3V0aWxzLnB5) | `22.50% <0.00%> (-65.00%)` | :arrow_down: |
| [pymc/printing.py](https://codecov.io/gh/pymc-devs/pymc/pull/5763/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9wcmludGluZy5weQ==) | `21.78% <0.00%> (-64.36%)` | :arrow_down: |
| [pymc/distributions/bound.py](https://codecov.io/gh/pymc-devs/pymc/pull/5763/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL2JvdW5kLnB5) | `45.54% <0.00%> (-54.46%)` | :arrow_down: |
| [pymc/distributions/censored.py](https://codecov.io/gh/pymc-devs/pymc/pull/5763/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL2NlbnNvcmVkLnB5) | `40.35% <0.00%> (-54.39%)` | :arrow_down: |
| [pymc/distributions/timeseries.py](https://codecov.io/gh/pymc-devs/pymc/pull/5763/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL3RpbWVzZXJpZXMucHk=) | `30.27% <0.00%> (-47.28%)` | :arrow_down: |
| [pymc/step\_methods/mlda.py](https://codecov.io/gh/pymc-devs/pymc/pull/5763/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGVwX21ldGhvZHMvbWxkYS5weQ==) | `50.41% <0.00%> (-45.97%)` | :arrow_down: |
| [pymc/step\_methods/hmc/quadpotential.py](https://codecov.io/gh/pymc-devs/pymc/pull/5763/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGVwX21ldGhvZHMvaG1jL3F1YWRwb3RlbnRpYWwucHk=) | `37.65% <0.00%> (-42.90%)` | :arrow_down: |
| [pymc/exceptions.py](https://codecov.io/gh/pymc-devs/pymc/pull/5763/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9leGNlcHRpb25zLnB5) | `58.62% <0.00%> (-41.38%)` | :arrow_down: |
| ... and [32 more](https://codecov.io/gh/pymc-devs/pymc/pull/5763/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | |

LukeLB: Failing test looks to be down to sampling error? 
```
test_sample_posterior_predictive_after_set_data_with_coord
E       AssertionError: 
E       Not equal to tolerance rtol=1e-07, atol=0.1
E       
E       Mismatched elements: 2 / 2 (100%)
E       Max absolute difference: 0.13669218
E       Max relative difference: 0.02331315
E        x: array([5, 6])
E        y: array([4.89599 , 5.863308]) 
```
And it's highlighted in an outstanding issue #5739.
LukeLB: This test is failing because y is not the correct shape, `y.eval().shape == (1, 1)`. Both feature and group are TensorConstants so the addition to add_coord is doing as intented with respect to type changes. I'm not sure I understand why this is still happening.....
```{python}
def test_add_coord_fixed_dims():
    """
    test that dims does not cause resizing of y. See #5181.
    """
    with pm.Model(coords={"feature": [1], "group": ["A", "B"]}):
        x = pm.Normal("x", 0, 1, dims="feature")
        y = pm.Normal("y", x[..., None], 1, dims=("feature", "group"))
    assert x.eval().shape == (1,)
    assert y.eval().shape == (1, 2)
```



michaelosthege: Maybe do more asserts?

```python
with pm.Model(coords={"feature": [1], "group": ["A", "B"]}):
    assert pmodel.dim_lengths["feature"].eval() == 1
    assert pmodel.dim_lengths["group"].eval() == 2
```
LukeLB: My thinking for that test was to make sure that size and dims can be passed at the same time and we get the expected shape out which will not have been possible before this type change in `add_coord()`. So maybe its worth having two tests, the one you've suggested, and the other we can use the case you highlight https://github.com/pymc-devs/pymc/issues/5181#issuecomment-970573037 to test that size and dims can be passed at the same time.  What do you think?

michaelosthege: When `size/shape` and `dims` are passed at the same time, the `size/shape` takes priority and `dims` are just tracked for labeling the dimensions in the final `xarray` objects.

That's why I suggested to _only_ pass dims, and actually it's a little concerning that the test doesn't pass.

Maybe call `add_coord` directly while passing `fixed`. For the variables depending on it there should be no difference between `aesara.shared()` or `at.constant()` dim lengths.

```python
with pm.Model() as pmodel:
    pmodel.add_coord("feature", [1], fixed=True/False)
    pmodel.add_coord("feature", ["A", "B"], fixed=True/False)
    assert pmodel.dim_lengths["feature"].eval() == 1
    assert pmodel.dim_lengths["group"].eval() == 2
    x = pm.Normal("x", 0, 1, dims="feature")
    y = pm.Normal("y", x[..., None], 1, dims=("feature", "group"))
assert x.eval().shape == (1,)
assert y.eval().shape == (1, 2)
```
michaelosthege: Oh, now I see why. The `x[..., None]` creates the RV as `(1, 1)` and because the dimensionality is as expected, there's no further resizing.

This works 👇 because then the RV is `(1,)` and repeated for the "group" dimension.

```python
    # ...
    x = pm.Normal("x", 0, 1, dims="feature")
    y = pm.Normal("y", x, 1, dims=("group", "feature"))
assert x.eval().shape == (1,)
assert y.eval().shape == (2, 1)
```
michaelosthege: Yes, that's already allowed, but where this test comes in is in combination with the changes from #5761.

Actually, it's a little unfortunate that these changes are on separate branches.
Your test case should (latest after adding a `pm.set_data`) close the coverage gaps from that other PR.

I'm inclined to merge #5761 even with those coverage gaps (I'll review it really carefully right now).
After that you could just `git rebase origin/main`, add two lines doing `pm.set_data` and assert the `ShapeError` or `ShapeWarning` and then we should be good to go 🤔
LukeLB: @michaelosthege OK I think I understand. I'll rebase and commit your suggested changes.
LukeLB: > ah damn, I just realized that we forgot to actually _use_ the new `fixed` kwarg.
> 
> I commented a change in one of the tests that should currently fail. Let me know if you want me to take over - you put in a lot of effort already and I don't want to get demotivated by all this back and forth

Not at all! The longer it goes the more I learn :) Thanks for the offer though.
LukeLB: Why does the assert statement on line 759 fail? Its created from coord so should be a `TensorConstant` but is `Subtensor{int64}.0`.
michaelosthege: > Why does the assert statement on line 759 fail? Its created from coord so should be a `TensorConstant` but is `Subtensor{int64}.0`.

The `Subtensor{int64}.0` is because from `pm.ConstantData([1,2,3], dims="foo")` it's effectively doing this: `dim_lengths["foo"] = at.constant([1,2,3]).shape[0]`.

So what needs to be done is a change in the logic behind `pm.ConstantData` (or rather `pm.Data`) to determine the constant data's shape non-symbolically first, and then call `add_coord(..., fixed=True)`.
LukeLB: OK so I think I have a solution which I have pushed but I'm running into a wall with a failing test which I can't seem to fix. Unfortunatley I won't be able to work on this till after the weekend so if there is a time pressure on this then it will be worth someone else picking it up. 
michaelosthege: Thanks @LukeLB, I will try to pick it up in the evening
canyon289: @LukeLB this weekend is fine! You've made great progress on this and don't want to take away joy of getting it merged!
michaelosthege: I rebased the branch and squashed the history into three logically separate commits:
1. ⏳ Removing the ShapeWarning in favor of raising earlier
2. ⏳ Adding optional `coords` to the `pm.Data` API, because we were taking coordinate values from the data variable, which does not work for >1 dimensions.
3. ⏳ Addition of the `fixed` kwarg (the actual goal of this PR)

I will force-push them one by one so we can confirm that the first two commits don't create new problems.
michaelosthege: @LukeLB @canyon289 @ricardoV94 please review and note the new changes I made in the `pm.Data` API
canyon289: Me too!

Patch:
diff --git a/pymc/data.py b/pymc/data.py
index 0f1a54d09..c03461cc5 100644
--- a/pymc/data.py
+++ b/pymc/data.py
@@ -20,7 +20,7 @@ import urllib.request
 import warnings
 
 from copy import copy
-from typing import Any, Dict, List, Optional, Sequence, Union, cast
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
 
 import aesara
 import aesara.tensor as at
@@ -466,9 +466,15 @@ def align_minibatches(batches=None):
                 rng.seed()
 
 
-def determine_coords(model, value, dims: Optional[Sequence[str]] = None) -> Dict[str, Sequence]:
+def determine_coords(
+    model,
+    value,
+    dims: Optional[Sequence[Optional[str]]] = None,
+    coords: Optional[Dict[str, Sequence]] = None,
+) -> Tuple[Dict[str, Sequence], Sequence[Optional[str]]]:
     """Determines coordinate values from data or the model (via ``dims``)."""
-    coords = {}
+    if coords is None:
+        coords = {}
 
     # If value is a df or a series, we interpret the index as coords:
     if hasattr(value, "index"):
@@ -499,10 +505,14 @@ def determine_coords(model, value, dims: Optional[Sequence[str]] = None) -> Dict
             )
         for size, dim in zip(value.shape, dims):
             coord = model.coords.get(dim, None)
-            if coord is None:
+            if coord is None and dim is not None:
                 coords[dim] = range(size)
 
-    return coords
+    if dims is None:
+        # TODO: Also determine dim names from the index
+        dims = [None] * np.ndim(value)
+
+    return coords, dims
 
 
 def ConstantData(
@@ -510,6 +520,7 @@ def ConstantData(
     value,
     *,
     dims: Optional[Sequence[str]] = None,
+    coords: Optional[Dict[str, Sequence]] = None,
     export_index_as_coords=False,
     **kwargs,
 ) -> TensorConstant:
@@ -522,6 +533,7 @@ def ConstantData(
         name,
         value,
         dims=dims,
+        coords=coords,
         export_index_as_coords=export_index_as_coords,
         mutable=False,
         **kwargs,
@@ -534,6 +546,7 @@ def MutableData(
     value,
     *,
     dims: Optional[Sequence[str]] = None,
+    coords: Optional[Dict[str, Sequence]] = None,
     export_index_as_coords=False,
     **kwargs,
 ) -> SharedVariable:
@@ -546,6 +559,7 @@ def MutableData(
         name,
         value,
         dims=dims,
+        coords=coords,
         export_index_as_coords=export_index_as_coords,
         mutable=True,
         **kwargs,
@@ -558,6 +572,7 @@ def Data(
     value,
     *,
     dims: Optional[Sequence[str]] = None,
+    coords: Optional[Dict[str, Sequence]] = None,
     export_index_as_coords=False,
     mutable: Optional[bool] = None,
     **kwargs,
@@ -588,9 +603,11 @@ def Data(
         :ref:`arviz:quickstart`.
         If this parameter is not specified, the random variables will not have dimension
         names.
+    coords : dict, optional
+        Coordinate values to set for new dimensions introduced by this ``Data`` variable.
     export_index_as_coords : bool, default=False
-        If True, the ``Data`` container will try to infer what the coordinates should be
-        if there is an index in ``value``.
+        If True, the ``Data`` container will try to infer what the coordinates
+        and dimension names should be if there is an index in ``value``.
     mutable : bool, optional
         Switches between creating a :class:`~aesara.compile.sharedvalue.SharedVariable`
         (``mutable=True``) vs. creating a :class:`~aesara.tensor.TensorConstant`
@@ -624,6 +641,9 @@ def Data(
     ...         model.set_data('data', data_vals)
     ...         idatas.append(pm.sample())
     """
+    if coords is None:
+        coords = {}
+
     if isinstance(value, list):
         value = np.array(value)
 
@@ -665,15 +685,28 @@ def Data(
             expected=x.ndim,
         )
 
-    coords = determine_coords(model, value, dims)
-
+    # Optionally infer coords and dims from the input value.
     if export_index_as_coords:
-        model.add_coords(coords)
-    elif dims:
+        coords, dims = determine_coords(model, value, dims)
+
+    if dims:
+        if not mutable:
+            # Use the dimension lengths from the before it was tensorified.
+            # These can still be tensors, but in many cases they are numeric.
+            xshape = np.shape(arr)
+        else:
+            xshape = x.shape
         # Register new dimension lengths
         for d, dname in enumerate(dims):
             if not dname in model.dim_lengths:
-                model.add_coord(dname, values=None, length=x.shape[d])
+                model.add_coord(
+                    name=dname,
+                    # 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],
+                )
 
     model.add_random_variable(x, dims=dims)
 
diff --git a/pymc/model.py b/pymc/model.py
index a3c20bb81..59ae618b8 100644
--- a/pymc/model.py
+++ b/pymc/model.py
@@ -61,7 +61,7 @@ from pymc.data import GenTensorVariable, Minibatch
 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, ShapeWarning
+from pymc.exceptions import ImputationWarning, SamplingError, ShapeError
 from pymc.initial_point import make_initial_point_fn
 from pymc.math import flatten_list
 from pymc.util import (
@@ -1071,8 +1071,9 @@ class Model(WithMemoization, metaclass=ContextMeta):
         self,
         name: str,
         values: Optional[Sequence] = None,
+        mutable: bool = False,
         *,
-        length: Optional[Variable] = None,
+        length: Optional[Union[int, Variable]] = None,
     ):
         """Registers a dimension coordinate with the model.
 
@@ -1084,9 +1085,12 @@ class Model(WithMemoization, metaclass=ContextMeta):
         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 symbolic scalar of the dimensions length.
-            Defaults to ``aesara.shared(len(values))``.
+            A scalar of the dimensions length.
+            Defaults to ``aesara.tensor.constant(len(values))``.
         """
         if name in {"draw", "chain", "__sample__"}:
             raise ValueError(
@@ -1097,7 +1101,9 @@ class Model(WithMemoization, metaclass=ContextMeta):
             raise ValueError(
                 f"Either `values` or `length` must be specified for the '{name}' dimension."
             )
-        if length is not None and not isinstance(length, Variable):
+        if isinstance(length, int):
+            length = at.constant(length)
+        elif length is not None and not isinstance(length, Variable):
             raise ValueError(
                 f"The `length` passed for the '{name}' coord must be an Aesara Variable or None."
             )
@@ -1109,14 +1115,17 @@ class Model(WithMemoization, metaclass=ContextMeta):
             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,
         coords: Dict[str, Optional[Sequence]],
         *,
-        lengths: Optional[Dict[str, Union[Variable, None]]] = None,
+        lengths: Optional[Dict[str, Optional[Union[int, Variable]]]] = None,
     ):
         """Vectorized version of ``Model.add_coord``."""
         if coords is None:
@@ -1180,24 +1189,20 @@ class Model(WithMemoization, metaclass=ContextMeta):
             # NOTE: If there are multiple pm.MutableData containers sharing this dim, but the user only
             #       changes the values for one of them, they will run into shape problems nonetheless.
             if length_changed:
+                if original_coords is not None:
+                    if new_coords is None:
+                        raise ValueError(
+                            f"The '{name}' variable already had {len(original_coords)} coord values defined for "
+                            f"its {dname} dimension. With the new values this dimension changes to length "
+                            f"{new_length}, so new coord values for the {dname} dimension are required."
+                        )
                 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`."
-                    )
-                if length_tensor.owner is None:
-                    # This is the case if the dimension was initialized
-                    # from custom coords, but dimension length was not
-                    # stored in TensorConstant e.g by 'fixed' set to False
-
-                    warnings.warn(
-                        f"You're changing the shape of a variable "
-                        f"in the '{dname}' dimension which was initialized "
-                        f"from coords. Make sure to update the corresponding "
-                        f"coords, otherwise you'll get shape issues.",
-                        ShapeWarning,
+                        "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]
@@ -1210,13 +1215,6 @@ class Model(WithMemoization, metaclass=ContextMeta):
                             actual=new_length,
                             expected=old_length,
                         )
-                if original_coords is not None:
-                    if new_coords is None:
-                        raise ValueError(
-                            f"The '{name}' variable already had {len(original_coords)} coord values defined for "
-                            f"its {dname} dimension. With the new values this dimension changes to length "
-                            f"{new_length}, so new coord values for the {dname} dimension are required."
-                        )
                 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)


Test patch:
diff --git a/pymc/tests/test_data_container.py b/pymc/tests/test_data_container.py
index 4398e9f6f..2983c26dd 100644
--- a/pymc/tests/test_data_container.py
+++ b/pymc/tests/test_data_container.py
@@ -19,7 +19,7 @@ import pytest
 
 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,29 +315,47 @@ class TestData(SeededTest):
         }
         # 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):
+        with pm.Model() as pmodel:
+            pm.ConstantData(
+                "population", [100, 200], dims="city", coords={"city": ["Tinyvil", "Minitown"]}
+            )
+            pm.MutableData(
+                "temperature",
+                [[15, 20, 22, 17], [18, 22, 21, 12]],
+                dims=("city", "season"),
+                coords={"season": ["winter", "spring", "summer", "fall"]},
+            )
+        assert "city" in pmodel.coords
+        assert "season" in pmodel.coords
+        assert pmodel.coords["city"] == ("Tinyvil", "Minitown")
+        assert pmodel.coords["season"] == ("winter", "spring", "summer", "fall")
+
     def test_symbolic_coords(self):
         """
         In v4 dimensions can be created without passing coordinate values.
         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
diff --git a/pymc/tests/test_model.py b/pymc/tests/test_model.py
index 50b06e24b..6a01bc28d 100644
--- a/pymc/tests/test_model.py
+++ b/pymc/tests/test_model.py
@@ -27,7 +27,9 @@ import pytest
 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
@@ -702,6 +704,67 @@ def test_nested_model_coords():
     assert set(m2.RV_dims) < set(m1.RV_dims)
 
 
+def test_shapeerror_from_resize_immutable_dims():
+    """
+    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")
+
+        m = 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]]})
+
+    with pytest.raises(ShapeError, match="was initialized from 'a'"):
+        # Can't work because the "fixed" dimension is linked to a constant shape:
+        # Note that the new data tries to change both dimensions
+        with pmodel:
+            pm.set_data({"m": [[1, 2], [3, 4]]})
+
+
+def test_valueerror_from_resize_without_coords_update():
+    """
+    Resizing a mutable dimension that had coords,
+    without passing new coords raises a ValueError.
+    """
+    with pm.Model() as pmodel:
+        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: