Consider making RandomVariable multivariate size under broadcasting more consistent

[brandonwillard]

Aesara's multivariate RandomVariables allow parameter broadcasting, while the NumPy equivalents do not. Use of the size parameter in this situation requires a choice in the way the parameter is interpreted. The current choice doesn't match the case of univariate broadcasting under size as well as it could, which can lead to confusion—or at least a less succinct rule for the resulting shape.

To illustrate, the following fails in NumPy

import numpy as np


>>> np.random.multivariate_normal(np.zeros((1, 2)), np.ones((1, 2, 2)), size=(3, 1))
...
ValueError: mean must be 1 dimensional

while the equivalent expression is supported in Aesara:

from aesara.tensor.random.basic import multivariate_normal


>>> multivariate_normal(np.zeros((1, 2)), np.ones((1, 2, 2)), size=(3,)).eval()
array([[[-0.67949394, -0.67949394]],

       [[ 0.58113935,  0.58113935]],

       [[ 0.17469251,  0.17469251]]])

Aesara interprets size=(3,) to mean "sample the distribution(s) with parameters np.zeros((1, 2)) and np.ones((1, 2, 2)) three times"; however, this clashes with NumPy's behavior for univariate distributions, which expects the size parameter to reflect both the additional dimensions (e.g. the 3 we want added) and the extra "broadcasted" dimensions implied by the parameters.

For instance, if we want three samples of two normal distributions with different mean parameters in NumPy, we use something like the following:

>>> np.random.normal(np.array([1, 10]), 1, size=(3, 2))
array([[ 1.16253515, 11.75758947],
       [ 1.29887451, 10.59990073],
       [-1.02698362,  9.41536206]])

If we set size=(3,) it would fail with a shape mismatch.

The same is true in Aesara for the univariate case, so, to make Aesara's additional multivariate broadcasting functionality more consistent with NumPy's univariate broadcasting, we could require that size be (3, 1) in our multivariate_normal example above. In other words, we can require size to include the extra broadcasted dimension—but not the "support" dimensions (i.e. (2,) and (2, 2) for the mean and covariance, respectively).

We could make such a change without breaking from NumPy's design, because this particular case isn't yet supported by NumPy.

This change would—however—make the use of multivariate distributions under broadcasting a little more cumbersome, because it would require one to unnecessarily specify an extra dimension. The alternative is to completely break with NumPy and change size to means something altogether different (e.g. the shape to which you want to duplicate the underlying broadcasted distribution, or append duplicates of the broadcasted distribution, etc.).

NB: Just to be clear, it's the multivariate_normal operation that's being broadcasted across the parameters, and not that the parameters are necessarily broadcasted—although they will be when they don't already match across dimensions.

[twiecki]

I think it's fair to say that numpy does not handle size very consistently between univariate and multivariate distributions. It's understandable as it's probably not very important to them. So in this instance I agree that we do not need to tie ourselves to their standard.

My vote goes to consistency and have size only ignore support dimensions.

CC @michaelosthege @ricardoV94

[ricardoV94]

This might provide some clues as to where scipy/ numpy might be going: numpy/numpy#17669

It seems there is an interest in extending the behavior over there

[brandonwillard]

That discussion is spot on; we really should follow it.

[brandonwillard]

Don't forget, as a step toward simplification without deviating from NumPy/SciPy, we can always implement an alternative keyword option—call it reps—that works in the desired way: i.e. by only creating additional "replicated" dimensions that extend the "base" independent and support dimensions determined by the parameters and underlying distributions.

[ricardoV94]

That sounds like the most intuitive parametrization (user facing-wise) even if under the hood we keep using size everywhere

[ricardoV94]

Coming back to this. What would we be using under the hood (e.g. in the rng_fn)?

When is the NumPy definition of size useful at all? Does it makes parameter broadcasting easier, or solve any ambiguities?

[brandonwillard]

Just so everyone knows, everything for this topic has been implemented in #446 (not the reps thing, though). We can simply rebase and merge that; otherwise, there might be some follow-ups necessary in AePPL and PyMC, but nothing big.