Multivariate Normal with Block Diagonal Covariance / Batched Multivariate Normal

Misc.
1

Dear Pyro-Forum,

I am trying to implement a multivariate Normal prior over columns of a matrix with different covariance matrices for each column, i.e. for matrix A we have a_.,j as individual columns of A and a_.,j ~ N(0, Sigma_j), for j=1,…,R.

I have tried a sequential plate statement:

core_a = torch.empty(self.output_dim, self.rank)
core_b = torch.empty(self.feature_dim, self.rank)
for j in pyro.plate(name='rank_dim', size=self.rank):
    core_a[:, j] = pyro.sample(name='column_a_{}'.format(j),
                         fn=dist.MultivariateNormal(loc=torch.zeros(self.output_dim),
                         covariance_matrix=cov_core_a[:, :, j]))

However, the AutoGuide I am currently using does not support sequential plates.

Does anyone have a suggestion how to vectorize this kind of prior?

Best Regrads!

2

you should just be able to use a ‘batched’ MultivariateNormal distribution. something like

cov = torch.zeros(K, D, D)
cov[0] = torch.eye(D)
cov[1] = 2.0 * torch.eye(D)
...
dist.MultivariateNormal(loc=torch.zeros(K, D), covariance_matrix=cov)

that way all computations will be vectorized

2 Likes
3

Dear @martinjankowiak,
thank you for the swift reply. It worked perfectly.