Implementing Multitask Fixed-Noise Gaussian Likelihood

[awang124]

The GPyTorch documentation gives extensive introduction and examples for defining models based on gpytorch.models.ExactGP. However, I can't any documentation / guide for defining custom likelihoods. For example, I aim to define a multitask version of gpytorch.likelihoods.FixedNoiseGaussianLikelihood, where, instead of passing a Tensor of shape (N,), with each datapoint's noise variance, I'd pass a Tensor of shape (N,T,T), with each point's full cross-task noise covariance matrix. How would I go about defining such a likelihood? And more importantly, how do I define my own likelihoods in general? What base class should I inherit from? What methods (i.e. forward, _shaped_noise_covar, marginal, etc.) need be defined / overriden? I'm new to the GPR realm, and I'd greatly appreciate some guidance here.

Currently, my model is:

class GPROFE(gpytorch.models.ExactGP):
    def __init__(self, X, Y, likelihood):
        super(GPROFE, self).__init__(X, Y, likelihood)
        self.mean_module = gpytorch.means.MultitaskMean(
            gpytorch.means.ConstantMean(), num_tasks=2
        )
        self.covar_module = gpytorch.kernels.MultitaskKernel(
            gpytorch.kernels.RBFKernel(), num_tasks=2, rank=1
        )
    def forward(self, X):
        mean = self.mean_module(X)
        covar = self.covar_module(X)
        return gpytorch.distributions.MultitaskMultivariateNormal(mean, covar)

And my likelihood is:

class FixedLikelihood(gpytorch.likelihoods.MultitaskGaussianLikelihood):
    def __init__(self, num_tasks, noise_covar):
        super().__init__(num_tasks=num_tasks, has_task_noise=False)
        self.register_buffer("noise_covar", noise_covar)
        self.raw_noise.requires_grad_(False)

    def forward(self, funcdist, *params, **kwargs):
        mean, covar = funcdist.mean, funcdist.lazy_covariance_matrix
        if isinstance(covar, gpytorch.lazy.LazyEvaluatedKernelTensor):
            covar = covar.evaluate_kernel()
        if funcdist._interleaved:
            noise_covar = torch.block_diag(*self.noise_covar)
        else:
            n = noise_covar.shape[0] * self.num_tasks
            t = self.num_tasks
            noise_covar = torch.zeros((n * t, n * t))
            for i in range(t):
                for j in range(t):
                    a = slice(n * i, n * (i + 1))
                    b = slice(n * j, n * (j + 1))
                    noise_covar[a, b] = torch.diag(self.noise_covar[..., i, j])
        covar = covar + noise_covar
        return funcdist.__class__(mean, covar, interleaved=funcdist._interleaved)

But I believe forward isn't being called during training.