Batch Independent Multioutput GP shared Gram matrix

[alegresor]

The Batch independent multioutput GP tutorial shows how to to set different kernel parameters for each output. For $N$ inputs and $M$ outputs this costs $\mathcal{O}(MN^3)$ as we need to invert $M$ Gram matrices, each $N \times N$. Is there a way to use a shared kernel for all models and therefore only have one gram matrix to invert at cost $\mathcal{O}(N^3)$? I suspect there is a way to do this by setting the batch_shape parameter for the kernel and mean, but I cant figure it out. Any help would be greatly appreciated. Thanks!

[alegresor]

Looks like PR #2470 proposes a fix to this issue.

[gpleiss]

If you call .expand(M, N, N) on the output of the kernel matrix, GPyTorch should use a BatchRepeatLinearOperator under the hood, which will have the computational costs of a N x N matrix.

[alegresor]

The following code

import gpytorch 
import torch 

class IndepVecGPShared(gpytorch.models.ExactGP):
    def __init__(self, x, y, mean_module=None, covar_module=None, likelihood=None):
        assert x.ndim==2 and y.ndim==2 and len(x)==len(y)
        likelihood = likelihood if likelihood is not None else gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=y.size(1),noise_constraint=gpytorch.constraints.GreaterThan(1e-8))
        super().__init__(x,y,likelihood)
        self.d_in = x.size(1)
        self.d_out = y.size(1)
        self.mean_module = mean_module if mean_module is not None else gpytorch.means.LinearMean(input_size=self.d_in,batch_shape=torch.Size([self.d_out]))
        self.covar_module = covar_module if covar_module is not None else gpytorch.kernels.ScaleKernel(gpytorch.kernels.RQKernel(ard_num_dims=self.d_in))
    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        covar_x = covar_x.expand(self.d_out,x.size(0),x.size(0))
        return gpytorch.distributions.MultitaskMultivariateNormal.from_batch_mvn(gpytorch.distributions.MultivariateNormal(mean_x,covar_x))
    
if __name__ == "__main__":
    gp = IndepVecGPShared(x=torch.rand(3,4),y=torch.rand(3,5))
    y = gp.forward(torch.rand(10,4))

    gp = gp.eval()
    y = gp(torch.rand(10,4))
    print(y)

gives the error

RuntimeError: x1 and x2 were not broadcastable with kernel of batch_shape torch.Size([]). Got x1.shape = torch.Size([2, 13, 4]) and x2.shape = torch.Size([5, 13, 4])

Am I setting this up wrong?