Implement shape and dims for v4

[twiecki]

@brandonwillard is proposing to remove the current shape kwarg and replace it with size which works more similar to the way numpy handles it. The difference is that size only specifies the batch dimension, while the event dimension is taken from the distribution. For univariate distributions, there is no difference, but for multivariate, there would be.

coords=dict(town=ABC, year=[2020, 2021])
X = MvNormal(cov=eye(3), shape=(2, 3))
Y = MvNormal(cov=eye(3), dims=("town", "year"))
Z = MvNormal(cov=eye(3) , size=2)

All three would lead to the same shaped (2, 3) MvNormal.

From @brandonwillard: "in v4, all the broadcasting and shape calculations for the arguments are already worked out so we always know what the shape of a single sample is under the given parameters. size just says that you want size-many more of those samples. as a matter of fact, size can be symbolic in v4!"

The big downside is that we're breaking backwards compatibility in a major way. Shapes are some of the hardest things to get correctly and I'm very concerned that users might shy away from v4 because of that.

For now, @brandonwillard @michaelosthege and I agree that definitely we should start with keeping shape for compatibility. The question is whether we should ever remove it completely, which is where I'm having concerns. After talking @junpenglao who also shared these concerns.

[ricardoV94]

Would specifying shape just be a redundant way of specifying size (+ some informative message if the shape didn't make sense)?

[michaelosthege]

shape and dims are fully redundant to each other.

shape and size are only directly comparable given the other parameters such as cov=eye(3).

Because shape ≡ dims and I think we all agree that dims was a great improvement of our API, I don't think we can/need/should make the internal logic go away.
However, we might still want to nudge everybody away from it.

One idea that circulated in our chat discussion was to slightly change what shape does in order to make it a new feature:

If distribution parameters are symbolic, it's not immediately clear how the RV looks like:

x = MvNormal(cov=cov_prior)

Specifying shape makes the model code easier to read, but makes it very rigid w.r.t. out-of-sample prediction:

x = MvNormal(cov=cov_prior, shape=(3,))   # can no longer be broadcasted into a different size

It's the same as

x = MvNormal(cov=cov_prior)
aesara.tensor.specify_shape(x, (3,))    # this Op needs to become a function output to be effective!

We could implement shape such that the SpecifyShape Op is automatically added, creating a more shape-rigid model.
The plus point then is that inputs to specify_shape could even be symbolic, creating asserts such as: "This RVs 2nd dimension must be the same size as that dim from the other RV."

[OriolAbril]

I like the idea of distinguishing batch from event dimensions. This is currently one of the main pain points (for me) with ArviZ converters. The variables in posterior predictive, observed data and log likelihood have the same shape for univariate distributions, but the log likelihood has a subset of the shape for multivariate.

ArviZ currently does a quick and dirty check when retrieving the log likelihood (for all converters, as AFAIK none of the libraries properly supported distinguish on this) which is dropping the last dimensions when a shape mismatch is found to see if this solves things.

The example above would therefore break ArviZ when using custom dims because the event dim is not the last.

coords=dict(town=ABC, year=[2020, 2021])
X = MvNormal(cov=eye(3), shape=(2, 3))
Y = MvNormal(cov=eye(3), dims=("year", "town"))
Z = MvNormal(cov=eye(3) , size=2)

would work. If the variable Y were observed, it would currently error out at the idata conversion step due to the shape specification.

I don't really know if this would help (I guess we can internally make sure the event dims are the last if using size) nor how to make it better, I just wanted to share that as it is a new feature with close to no docs and literally no examples yet. Possible improvements on that for the v4 converter would be great

[ricardoV94]

Shape was definitely one of those "invisible" features that took me some time to figure out. Looking at questions in the discourse, it seems to crop up frequently. The other that seems to confuse people (although it's just theano borrowing from the numpy API) is fancy indexing:

indexes = np.array([0, 1, 2, 3, 0, 1, 2, 3])
a = pm.Normal('a', 0, 1, shape=4)
like = pm.Normaly('like', a[indexes], 1, observed=data)

The new coords seems even more enigmatic to me. Would be nice to have a doc notebook just walking through the different shape/size and indexing possibilities in pymc3. I know we have some examples, but I think (without proof) that they tend to be difficult to find.

Edit: This was brought up here #4007

[brandonwillard]

Shapes are some of the hardest things to get correctly and I'm very concerned that users might shy away from v4 because of that.

In v4 users are no longer required to determine the resulting shapes of the objects they create. In v3 a user is forced to do this themselves, and, as we all know, it's confusing and leads to numerous errors. Those are actual reasons for users to shy away from PyMC. I don't understand how asking a user to do less is off-putting.

Also, I don't see any real advantages to porting the shape keyword—aside from backward compatibility, but, since we're explicitly creating a new version, backward compatibility breaks like this are completely justified.

More importantly, the shape keyword has a few very real disadvantages, and we've been dealing with them for quite a while now (e.g. it's confusing, error prone, limits the flexibility and reuse of our models, etc.), so I'm very much against propagating this key issue into the next version.

As @michaelosthege mentioned, if one wants to constrain the shape of a variable, they can use specify_shape, and we could implement a shape keyword option using that; however, a user that very specifically wants to restrict the shape of a random variable can just as easily apply specify_shape after creating said variable.

The problem I see with advertising a shape keyword in v4 is that v3 users who have learned this bad habit will likely end up doing it in v4, giving rise to all those old issues and limitations. This is also likely to happen when new users attempt to recreate models written for v3.

This shape keyword is both the interface to—and cause of—one of the biggest problems we're fixing in v4, so it's the last thing we want to unnecessarily shoehorn back into it.

[twiecki]

I suppose I would need to be sold with examples where size is superior. If we can market this as a big improvement it's a different story.

…

On Sat, Mar 20, 2021, 00:07 Brandon T. Willard ***@***.***> wrote: Shapes are some of the hardest things to get correctly and I'm very concerned that users might shy away from v4 because of that. In v4 users are no longer required to determine the resulting shapes of the objects they create. In v3 a user is forced to do this themselves, and, as we all know, it's confusing and leads to numerous errors. *Those* are actual reasons for users to shy away from PyMC. I don't understand how asking a user to do *less* is off-putting. Also, I don't see any real *advantages* to porting the shape keyword—aside from backward compatibility, but, since we're explicitly creating a *new version*, backward compatibility breaks like this are completely justified. More importantly, the shape keyword has a few very real *disadvantages*, and we've been dealing with them for quite a while now (e.g. it's confusing, error prone, limits the flexibility and reuse of our models, etc.), so I'm very much against propagating this key issue into the next version. As @michaelosthege <https://github.com/michaelosthege> mentioned, if one wants to constrain the shape of a variable, they can use specify_shape, and we *could* implement a shape keyword option using that; however, a user that very specifically wants to restrict the shape of a random variable can just as easily apply specify_shape after creating said variable. The problem I see with advertising a shape keyword in v4 is that v3 users who have learned this bad habit will likely end up doing it in v4, giving rise to all those old issues and limitations. This is also likely to happen when new users attempt to recreate models written for v3. This shape keyword is both the interface to—and cause of—one of the biggest problems we're fixing in v4, so it's the last thing we want to unnecessarily shoehorn back into it. — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#4552 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAFETGEI3VAK6X3WJ7CDODTTEPKMZANCNFSM4ZKXGZTQ> .

[brandonwillard]

I suppose I would need to be sold with examples where size is superior.

Here's a very simple example that works in my local #4551 branch of v4:

import numpy as np
import pymc3 as pm

import aesara
import aesara.tensor as aet


data = np.random.normal(size=(1, 3))
X_at = aesara.shared(data)


with pm.Model() as m:
    beta = pm.MvNormal("beta", mu=aet.zeros((X_at.shape[1],)), cov=aet.eye(X_at.shape[1]))
    Y = pm.MvNormal("Y", mu=X_at.dot(beta), cov=aet.eye(X_at.shape[0]), size=(2,))

with m:
    priorp_samples = pm.sample_prior_predictive(samples=10)

>>> priorp_samples["beta"].shape
(10, 3)
>>> priorp_samples["Y"].shape
(10, 2, 1)

new_data = np.random.normal(size=(2, 4))
X_at.set_value(new_data)

with m:
    priorp_samples = pm.sample_prior_predictive(samples=10)

>>> priorp_samples["beta"].shape
(10, 4)
>>> priorp_samples["Y"].shape
(10, 2, 2)

This would not even be remotely possible if the shapes of beta and/or Y were fixed via a shape keyword.

Otherwise, I invite everyone to imagine how this implementation would change if symbolic values were to be used in the shape value. I guarantee that the change wouldn't be an improvement in any way.

[twiecki]

OK that's awesome, I get it now. Shapes are automatically propagated through.

…

On Sat, Mar 20, 2021, 01:24 Brandon T. Willard ***@***.***> wrote: I suppose I would need to be sold with examples where size is superior. If Here's a very simple example that works in my local #4551 <#4551> branch of v4: import numpy as npimport pymc3 as pm import aesaraimport aesara.tensor as aet data = np.random.normal(size=(1, 3))X_at = aesara.shared(data) with pm.Model() as m: beta = pm.MvNormal("beta", mu=aet.zeros((X_at.shape[1],)), cov=aet.eye(X_at.shape[1])) Y = pm.MvNormal("Y", mu=X_at.dot(beta), cov=aet.eye(X_at.shape[0]), size=(2,)) with m: priorp_samples = pm.sample_prior_predictive(samples=10) >>> priorp_samples["beta"].shape (10, 3)>>> priorp_samples["Y"].shape (10, 2, 2) new_data = np.random.normal(size=(2, 4))X_at.set_value(new_data) with m: priorp_samples = pm.sample_prior_predictive(samples=10) >>> priorp_samples["beta"].shape (10, 4)>>> priorp_samples["Y"].shape (10, 2, 2) This would not even be remotely possible if the shapes of beta and/or Y were fixed via a shape keyword. Otherwise, I invite everyone to imagine how this implementation would change if symbolic values were to be used in the shape value. I guarantee that the change wouldn't be an improvement in any way. — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#4552 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAFETGCBMNK2SSAX33TAAKDTEPTNFANCNFSM4ZKXGZTQ> .

[michaelosthege]

The fact that this use case is now possible is awesome!
At the same time it's not something we need from every model.

In 95 % of the models I write, I'm more concerned with the inference on the data I have, not the out-of-sample prediction using shared variables.
For inference the dims kwarg is incredibly useful. It just happens that it's quite similar to shape. Our discussion should be titled "dims and size" not "shape vs size".

[brandonwillard]

In 95 % of the models I write, I'm more concerned with the inference on the data I have, not the out-of-sample prediction using shared variables.

Don't forget that, because we can do this shared variable thing, it also means that people who run the same model on many different datasets will no longer need to recreate their Models and recompile their respective log-likelihood graphs to C/whatever else.

[michaelosthege]

Don't forget that, because we can do this shared variable thing, it also means that people who run the same model on many different datasets will no longer need to recreate their Models and recompile their respective log-likelihood graphs to C/whatever else.

Sure, that's great. Let's figure out how to make dims work with that.

1. The values passed via dims are already sort of "symbolic". They are placeholders. Can we translate them into symbolic elements for size?
2. What about those dims that refer to implied dimensions? (That's actually very common. e.g. https://github.com/rtcovidlive/rtlive-global/blob/master/rtlive/model.py#L168-L176)
3. What happens to coords when the value of a shared variable is changed? Can we pass new values for the coords via something like pm.Data().set_value(val, coords=...)?

[OriolAbril]

From a naive/conceptual point of view, I see dims fitting better the size approach. Currently dims takes the info of coords to get the length of the dimension, but I actually don't think this is ideal.

Whenever I have to do cross-validation, I often use shape to define the model and then when converting to inferencedata, I use dims but no coords. I really don't care about coordinate values, and having to add extra code to sort coord values at model creation is work I don't really need and won't use, whereas having the correct dims in place helps with automatic alignment and broadcasting when postprocessing with xarray.

It would be great to pair named dimensions with the size argument and then define the coordinate values at sampling/conversion time which is what feels more natural. In this crossvalidation/out of sample sampling, the dimension id generally is part of the model and I think it would be useful to allow having this info still there like we do now with dims. The coordinate values however, are not part of the model and I think they should be better off not being tied to a model but to data/sampling realizations

[twiecki]

The benefit I see is the shape propagation, not necessarily that we can switch things for prediction. With shape or dim I always have to specify the full shape of every tensor, while with size I just have to remember the dimension I want to batch over, which is much simpler. I think dims could be changed in this way where I'm just specifying which variable I want to batch over. I changed the example below: we have two experiments/replications on a drug with different dose levels administered to multiple subjects.

import numpy as np
import pymc3 as pm

import aesara
import aesara.tensor as aet

n_doses = 3
n_subjects = 5

data = np.random.normal(size=(n_doses, n_subjects))
X_at = aesara.shared(data)


with pm.Model() as m:
    beta = pm.MvNormal(
        "beta", 
        mu=aet.zeros((n_doses,)), 
        cov=aet.eye(n_doses))

    Y = pm.MvNormal(
        "Y", 
        mu=X_at.dot(beta), 
        cov=aet.eye(n_states), 
        dims=("experiments"))

Without this, the likelihood would have to look like:

    Y = pm.MvNormal(
        "Y", 
        mu=X_at.dot(beta), 
        cov=aet.eye(n_states), 
        dims=("experiments", "doses", "subjects"))

Maybe we want to rename dim in the above case to plate_dim, batch_dim, or rep_dim to avoid confusion.