Nested models don't inherit their parent model's coords

[lucianopaz]

Description of your problem

Based off of this discussion, nested models should share their coordinates, but they don't at the moment, which leads to this error

COORDS = {"dim": range(10)}
with pm.Model(name="m1", coords=COORDS) as m1:
    a = pm.Normal("a")
    with pm.Model(name="m2") as m2:
        b = pm.Normal("b")
        c = pm.HalfNormal("c")
        d = pm.Normal("d", b, c, dims="dim")
    e = pm.Normal("e", a + d, dims="dim")

Please provide the full traceback.

KeyError: "Dimensions {'dim'} are unknown to the model and cannot be used to specify a `size`."

KeyError                                  Traceback (most recent call last)
<ipython-input-2-c89cca9d3d36> in <module>
      5         b = pm.Normal("b")
      6         c = pm.HalfNormal("c")
----> 7         d = pm.Normal("d", b, c, dims="dim")
      8     e = pm.Normal("e", a + d, dims="dim")
      9 

~/repos/pymc3/pymc/distributions/distribution.py in __new__(cls, name, rng, dims, initval, observed, total_size, transform, *args, **kwargs)
    251         # Create the RV and process dims and observed to determine
    252         # a shape by which the created RV may need to be resized.
--> 253         rv_out, dims, observed, resize_shape = _make_rv_and_resize_shape(
    254             cls=cls, dims=dims, model=model, observed=observed, args=args, rng=rng, **kwargs
    255         )

~/repos/pymc3/pymc/distributions/distribution.py in _make_rv_and_resize_shape(cls, dims, model, observed, args, **kwargs)
    160     if dims is not None:
    161         if dims_can_resize:
--> 162             resize_shape, dims = resize_from_dims(dims, ndim_actual, model)
    163         elif Ellipsis in dims:
    164             # Replace ... with None entries to match the actual dimensionality.

~/repos/pymc3/pymc/distributions/shape_utils.py in resize_from_dims(dims, ndim_implied, model)
    518     unknowndim_resize_dims = set(dims[:ndim_resize]) - set(model.dim_lengths)
    519     if unknowndim_resize_dims:
--> 520         raise KeyError(
    521             f"Dimensions {unknowndim_resize_dims} are unknown to the model and cannot be used to specify a `size`."
    522         )

KeyError: "Dimensions {'dim'} are unknown to the model and cannot be used to specify a `size`."

Please provide any additional information below.

Versions and main components

• PyMC/PyMC3 Version: 4.0.0b1
• Aesara/Theano Version: 2.3.2
• Python Version: 3.9.7
• Operating system: Mac
• How did you install PyMC/PyMC3: (conda/pip) pip

[ricardoV94]

@lucianopaz can I assign this issue to you, or should we ask for help?