Is it possible to declare multiple MvNormal variables with different covariance matrices in a vectorized way?

Questions v5
4

I wrote this code for doing it a while ago, but maybe kroneker products is a better choice, like this:

# Make a set of (N, N) matrices of all zeros with a 1 in the (i,i)th position
indicators = [np.zeros((5, 5)) for _ in range(5)]
for i in range(5):
    indicators[i][i, i] = 1

with pm.Model() as mod:
    L = pm.Normal('L', size=(5, 3, 3))
    # This is just to make a block of positive semi-definite matrices, you can ignore it
    # pt.einsum when?
    covs, _ = pytensor.scan(lambda x: x @ x.T, sequences=[L])
     
    # Here's where the actual block diagonal matrix get made
    big_cov2 = pt.stack([pt.linalg.kron(indicators[i], covs[i]) for i in range(5)]).sum(axis=0)
    obs = pm.MvNormal('obs', mu=0, cov=big_cov)

No idea which is faster, but this way is certainly less code, and avoids a scan.