Pymc3: Bug in MatrixNormal.sample

[aseyboldt]

If one of the matrix parameters is a random variable, then MatrixNormal.sample runs into shape errors because it might be passed an array of col or row matrices instead of just one. pm.sample_prior_predictive will fail in a case like this:

import numpy as np
import pymc3 as pm
K = 3
D = 15
mu_0 = np.zeros((D,K))
lambd = 1.0
with pm.Model() as model:
    sd_dist = pm.HalfCauchy.dist(beta=2.5)
    packedL = pm.LKJCholeskyCov(f"packedL",eta=2, n=D, sd_dist=sd_dist)
    L = pm.expand_packed_triangular(D, packedL, lower=True)
    Sigma = pm.Deterministic(f"Sigma", L.dot(L.T)) # D x D covariance
    mu = pm.MatrixNormal(
        f"mu", 
        mu=mu_0, 
        rowcov=(1 / lambd) * Sigma, 
        colcov = np.eye(K), shape=(D,K))
    prior = pm.sample_prior_predictive(2)

In this example rowchol will have shape (2, 15, 15).
https://github.com/pymc-devs/pymc3/blob/master/pymc3/distributions/multivariate.py#L1598

Originally reported by calvinm on discord:
https://discourse.pymc.io/t/reshape-error-inside-sample-prior-predictive/3715

[rpgoldman]

Should we add this to the tests, as an expected fail for now?

[aseyboldt]

Maybe it is not that hard to fix, we could just broadcast both matrices and mu to (sizeprod, M, N), and after making the samples reshape back to size + (M, N)...

[lucianopaz]

This issue is similar to #2848. Most, if not all, of the multivariate distributions have some shape handling defects when their distribution parameters depend on other RVs.

[michael-ziedalski]

I am just wondering now then, what would hypothetically need to be done, as a general fix, to get rid of this bug for the rest of the distributions?

[Sayam753]

I am just wondering now then, what would hypothetically need to be done, as a general fix, to get rid of this bug for the rest of the distributions?

I sense broadcasting all the parameters just after they are drawn and then proceeding with calculations, will be a good option. That's how MvNormal.random method is fixed in #4207.