I have been trying to use GPytorch do Hadamard Multitask GP Regression, for batched input data. Copying the tutorial exactly allows me to correctly perform the regression on a single vector input x.

I cannot get it to work for batched inputs of the shape [batch, num_datapoints, dim_x]. I have taken a look at this tutorial, but still couldn’t get it to work.

My current code is the following:

likelihood = gpytorch.likelihoods.GaussianLikelihood(batch_shape=xc.shape[:1])

xc_labels = torch.concat((
    torch.zeros(10, 1), # dummy values for example
    torch.ones(10, 1), 
    dim=0
).to(torch.long)

model = MultitaskGPModel((xc, xc_labels), yc, self.kernel, likelihood, out_dim=2)
model.eval()

outputDist = likelihood(model(xt, torch.zeros(xt.shape[-2], 1 , dtype=torch.long)))

class MultitaskGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, kernel, likelihood, out_dim):
        super(MultitaskGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ZeroMean(batch_shape=train_x[0].shape[:1])
        self.covar_module = kernel
        self.covar_module.batch_shape = train_x[0].shape[:1]
        self.task_module = gpytorch.kernels.IndexKernel(num_tasks=out_dim)

    def forward(self, x, i):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        covar_i = self.task_module(i)
        covar = covar_x.mul(covar_i)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar)

Where kernel is a simple RBF kernel and xc is of shape [batch, num_data, x_dim] and yc of shape [batch,num_data, 1]. This fails with the error:

Flattening the training labels failed. The most common cause of this
error is that the shapes of the prior mean and the training labels are
mismatched. The shape of the train targets is torch.Size([1, 110, 1]),
while the reported shape of the mean is torch.Size([1, 110])

If I instead use an xc of shape [num_data] (i.e. no batches and 1-dimensional x_dim) and do not set the batch_shape anywhere it works just fine, but I really want to be able to do this for batched inputs.

How can I properly do a batched Hadamard Multitask GP regression?