Add a shape check to MvNormal when cov and mean are of different length

[Sayam753]

Description of your problem or feature request

When mean and cov parameters for MvNormal distribution are not of same length, a long traceback is shown. A shape check can be added to improve the error message, to replicate NumPy's behaviour.

Please provide a minimal, self-contained, and reproducible example.

>>> import numpy as np
>>> np.__version__
'1.21.2'
>>> rng = np.random.default_rng()
>>> rng.multivariate_normal(mean=np.ones(1), cov=np.eye(3))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "_generator.pyx", line 3590, in numpy.random._generator.Generator.multivariate_normal
ValueError: mean and cov must have same length

>>> import aesara
>>> aesara.__version__
'2.2.6'
>>> from aesara.tensor.random.basic import multivariate_normal
>>> multivariate_normal(mean=np.zeros(1), cov=np.eye(3)).eval()

Please provide the full traceback of any errors.

Traceback

Traceback (most recent call last):
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/compile/function/types.py", line 977, in __call__
    if output_subset is None
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/graph/op.py", line 500, in rval
    r = p(n, [x[0] for x in i], o)
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/tensor/random/op.py", line 392, in perform
    smpl_val = self.rng_fn(rng, *(args + [size]))
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/tensor/random/basic.py", line 348, in rng_fn
    return safe_multivariate_normal(mean, cov, size=size, rng=rng)
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/tensor/random/basic.py", line 299, in safe_multivariate_normal
    stats.multivariate_normal(mean=mean, cov=cov, allow_singular=True).rvs(
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/scipy/stats/_multivariate.py", line 362, in __call__
    seed=seed)
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/scipy/stats/_multivariate.py", line 729, in __init__
    None, mean, cov)
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/scipy/stats/_multivariate.py", line 400, in _process_parameters
    cov.shape = (1, 1)
ValueError: cannot reshape array of size 9 into shape (1,1)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/graph/basic.py", line 553, in eval
    rval = self._fn_cache[inputs](*args)
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/compile/function/types.py", line 993, in __call__
    storage_map=getattr(self.fn, "storage_map", None),
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/link/utils.py", line 526, in raise_with_op
    raise exc_value.with_traceback(exc_trace)
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/compile/function/types.py", line 977, in __call__
    if output_subset is None
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/graph/op.py", line 500, in rval
    r = p(n, [x[0] for x in i], o)
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/tensor/random/op.py", line 392, in perform
    smpl_val = self.rng_fn(rng, *(args + [size]))
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/tensor/random/basic.py", line 348, in rng_fn
    return safe_multivariate_normal(mean, cov, size=size, rng=rng)
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/aesara/tensor/random/basic.py", line 299, in safe_multivariate_normal
    stats.multivariate_normal(mean=mean, cov=cov, allow_singular=True).rvs(
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/scipy/stats/_multivariate.py", line 362, in __call__
    seed=seed)
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/scipy/stats/_multivariate.py", line 729, in __init__
    None, mean, cov)
  File "/Users/user/miniconda3/envs/pymc_v4/lib/python3.7/site-packages/scipy/stats/_multivariate.py", line 400, in _process_parameters
    cov.shape = (1, 1)
ValueError: cannot reshape array of size 9 into shape (1,1)
Apply node that caused the error: multivariate_normal_rv{1, (1, 2), floatX, False}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FD24FC89230>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{(1,) of 0.0}, TensorConstant{[[1. 0. 0...0. 0. 1.]]})
Toposort index: 0
Inputs types: [RandomGeneratorType, TensorType(int64, vector), TensorType(int64, scalar), TensorType(float64, (True,)), TensorType(float64, matrix)]
Inputs shapes: ['No shapes', (0,), (), (1,), (3, 3)]
Inputs strides: ['No strides', (8,), (), (8,), (24, 8)]
Inputs values: [Generator(PCG64) at 0x7FD24FC89230, array([], dtype=int64), array(11), array([0.]), 'not shown']
Outputs clients: [[], ['output']]

Backtrace when the node is created (use Aesara flag traceback__limit=N to make it longer):
  File "<stdin>", line 1, in <module>

HINT: Use the Aesara flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

Versions and main components

• Python version: 3.7.7
• Operating system: Mac OS 11.6
• How did you install Aesara: pip

[brandonwillard]

The multivariate normal Op in Aesara uses scipy.stats.multivariate_normal and not NumPy's, so that's one reason you're seeing different error messages.

[ricardoV94]

The shape inference in this case is also off (assuming this should work at all):

at.random.multivariate_normal(mean=np.zeros(1), cov=np.eye(3)).shape.eval()
# array([1])

Or is the shape not supposed to worry about valid parameters? Edit: it seems that is indeed the case: https://aesara.readthedocs.io/en/latest/tutorial/shape_info.html#shape-inference-problem

[brandonwillard]

Or is the shape not supposed to worry about valid parameters?

Remember, some things—like parameter validation and shape checking—can only be adequately performed at "run-time".

These types of questions/requests appear a lot and tend to involve that symbolic vs. run-time confusion, especially when Constants are involved in the examples.

It is technically possible for some extra checking to occur at graph construction when the inputs are constants, but, in most cases, it would involve a lot of extra, specialized logic that would only provide a somewhat more convenient error in special cases.

[Sayam753]

But, in the case -

at.random.multivariate_normal(mean=np.zeros(1), cov=np.eye(3)).shape.eval()

why cannot we broadcast mean to shape (3)?

[brandonwillard]

why cannot we broadcast mean to shape (3)?

The Op's broadcast pattern/signature is (n), (n, n) -> (n), but the example's arguments imply (1), (3, 3) -> (?).

[Sayam753]

I meant to say, for the case

at.random.multivariate_normal(mean=np.zeros(1), cov=np.eye(3)).shape.eval()

arguments imply (1), (3, 3) -> (3), (3, 3) -> (3) i.e adding an intermediate step to match the Op's broadcast pattern.

[brandonwillard]

I meant to say, for the case

at.random.multivariate_normal(mean=np.zeros(1), cov=np.eye(3)).shape.eval()

arguments imply (1), (3, 3) -> (3), (3, 3) -> (3) i.e adding an intermediate step to match the Op's broadcast pattern.

The broadcast pattern must match the "base" input and output dimensions of the mathematical object being modeled (i.e. a multivariate normal random variable under a specific parameterization), and that doesn't.

[ricardoV94]

To give some context, numpy vectorize would raise a ValueError in this case:

import numpy as np

def func(mu, sigma):
    return mu

vfunc = np.vectorize(func, signature='(n),(n,n)->(n)')
vfunc(np.ones(1), np.eye(3))

ValueError: inconsistent size for core dimension 'n': 3 vs 1

Which is kind of what we are doing with RandomVariables