RV size is not correctly inferred from dims

[lucianopaz]

Description of your problem

While writing a hierarchical model, I ran into a shape problem. Minimally, what I want to do looks like this:

>>> 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"))
... y.shape.eval()
WARNING (aesara.tensor.basic_opt): Failed to infer_shape from Op normal_rv{0, (0, 0), floatX, True}.
Input shapes: [None, (TensorConstant{0},), (), (<TensorType(int64, scalar)>, TensorConstant{1}), ()]
Exception encountered during infer_shape: <class 'Exception'>
Exception message: sharedvar.value does not exist anymore. Use sharedvar.get_value() or sharedvar.set_value() instead.
Traceback: Traceback (most recent call last):
  File "/home/lpaz/repos/aesara/aesara/tensor/basic_opt.py", line 867, in get_node_infer_shape
    o_shapes = shape_infer(
  File "/home/lpaz/repos/aesara/aesara/tensor/random/op.py", line 299, in infer_shape
    shape = self._infer_shape(size, dist_params, param_shapes=param_shapes)
  File "/home/lpaz/repos/aesara/aesara/tensor/random/op.py", line 196, in _infer_shape
    param_shapes = params_broadcast_shapes(
  File "/home/lpaz/repos/aesara/aesara/tensor/random/utils.py", line 47, in params_broadcast_shapes
    rev_extra_dims = [
  File "/home/lpaz/repos/aesara/aesara/tensor/random/utils.py", line 48, in <listcomp>
    max_bcast(a, b)
  File "/home/lpaz/repos/aesara/aesara/tensor/random/utils.py", line 41, in max_bcast
    if getattr(x, "value", x) == 1:
  File "/home/lpaz/repos/aesara/aesara/compile/sharedvalue.py", line 211, in _value_get
    raise Exception(
Exception: sharedvar.value does not exist anymore. Use sharedvar.get_value() or sharedvar.set_value() instead.
array([1, 1])

So the shape of y is inferred incorrectly, and this led to problems downstream in the full model. If instead of using dims, I use size, everything works fine.

>>> with pm.Model(coords={"feature": [1], "group": ["A", "B"]}):
...     x = pm.Normal("x", 0, 1, size=(1,))
...     y = pm.Normal("y", x[..., None], 1, size=(1, 2))
... y.shape.eval()
array([1, 2])

Versions and main components

• PyMC/PyMC3 Version: main branch at commit 073e26b
• Aesara/Theano Version: 2.2.6
• Python Version: 3.8.5
• Operating system: Ubuntu 18.04
• How did you install PyMC/PyMC3: pip

[michaelosthege]

There are two levels to this problem.

First, your x[..., None].shape.eval() already has 2 dimensions and since this helper function only looks at ndim no resizing is triggered in Distribution.__new__.
We could only do that resize if we were to evaluate and compare the shape of the incoming variable and the symbolic dim_length stored in the model.

The second level is that new dimensions are always added to the left of the shape tuple.
So if x was a ndim=1 column vector, you'd need to specify dims=("group", "feature"), or dims=("group", ...) and then it works:

with pm.Model(coords={
    "feature": [1],
    "group": ["A", "B"]
}) as pmodel:
    x = pm.Normal("x", 0, 1, dims=("feature",))
    y = pm.Normal("y", x, 1, dims=("group", ...))

y.shape.eval()
# >>> array([2, 1], dtype=int64)

Note that when using the Ellipsis feature, the corresponding dimensions will remain unnamed. That's because dim names are not inferred through the Aesara graph.

[michaelosthege]

I forgot to crosslink: #4656

It shouldn't be too difficult to realize and then the following should work too:

with pm.Model(coords={"feature": [1], "group": ["A", "B"]}):
    x = pm.Normal("x", 0, 1, dims="feature")
    y = pm.Normal("y", x[..., None], 1, size=(1, 2), dims=("feature", "group"))
    # simultaneous size & dims 👆.
    # 👉 size causes the resize you desired
    # 👉 dims have no effect on the shape, but wll end up in the InferenceData

[michaelosthege]

@lucianopaz are you satisfied with the explanations above? Can we close this issue?

[lucianopaz]

@lucianopaz are you satisfied with the explanations above? Can we close this issue?

Adding to the left of the core dimensions implied by the parameter does work. But I find it odd that size and dims behave differently. Maybe we should add some extra functionality to resize_from_dims to get this consistent?

[michaelosthege]

I fear added functionality of resize_from_dims can end us up with conflicting scenarios, particularly as soon as we allow to pass dims in addition to shape/size (#4656).

Instead I would propose the following:

1. Allow for passing size/shape and dims (Allow naming of dims when parametrizing with size #4656) since that's an often seen use case (e.g. every time someone is doing sth like pm.Normal("n", 1, 2, size=5, dims="year") where the "year" dim is previously undefined).
2. With that we can have the clear rule "size/shape take priority over dims" and decide if we want to rv.shape.eval() automatically to detect shape/dims mismatches right at RV creation time.

About 2.) I'm a little sceptical because of runtime overhead, NotImplementedError or rng random states, but it may still be worth it, given that before (now and v3) incorrect dims only blew up after sampling.

[michaelosthege]

I think the core issue here is now captured by #5613 ?

And specifically by this comment.