Multitask GP syntactic sugar

Multitask GPs can be written in GPyTorch as is (see #208), but right now it is extremely ugly.

Proposed interface:

• Introduce a MultitaskKernel model, which wraps a standard kernel function and returns the Kronecker product between an index (task) kernel and the supplied input kernel.
• Mean functions take in an n_tasks argument, which defaults as None. If n_tasks=d, means will return an n x d output.
• Multiple outputs (for train_y) should be expressed in an n x d matrix.

Necessary updates:

• We introduce a syntactic sugar for GaussianRandomVariables. If the mean is an n x d matrix and the covariance is an nd x nd LazyVariable, we will figure things out
  • var, std, etc. will return n x d matrices
  • sample will return an n x d x s tensor
  • Explicitly evaluating the covariance might return an n x d x n x d tensor, or maybe a nd x nd tensor? Thoughts? (Probably won’t happen too often, but probably a good corner case to catch.)
• inv_quad_log_det is updated to handle this new syntactic sugar.

The hadamard multitask case (one task per input) will remain essentially the same.

cc @Balandat @darbour @jacobrgardner

[Balandat]

Here is something quick an dirty that takes the n x t mean and nt x nt covar and gets the job done at least for sampling:

class MultiOutputGaussianRandomVariable(RandomVariable):
    def __init__(self, mean, covar):
        """
        Constructs a multi-output multivariate Gaussian random variable, based on mean and covariance
        Can be multi-output multivariate, or a batch of multi-output multivariate Gaussians

        Passing a matrix mean corresponds to a multi-output multivariate Gaussian
        Passing a matrix mean corresponds to a batch of multivariate Gaussians

        Params:
        - mean (Tensor: matrix n x t, or batch matrix b x n x t) mean of Gaussian distribution
        - covar (Tensor: matrix nt x nt, or batch matrix b x nt x nt) covariance of Gaussian distribution
        """
        super(MultiOutputGaussianRandomVariable, self).__init__(mean, covar)
        if not torch.is_tensor(mean) and not isinstance(mean, LazyVariable):
            raise RuntimeError("The mean of a GaussianRandomVariable must be a Tensor or LazyVariable")

        if not torch.is_tensor(covar) and not isinstance(covar, LazyVariable):
            raise RuntimeError("The covariance of a GaussianRandomVariable must be a Tensor or LazyVariable")
            
        if mean.ndimension() not in {2, 3}:
            raise RuntimeError("mean should be a matrix or a batch matrix (batch mode)")

        if not isinstance(covar, LazyVariable):
            covar = NonLazyVariable(covar)
        
        self._mean = mean
        self._covar = covar
        self._nbatch = mean.shape[0] if mean.ndimension() == 3 else None
        self._n = mean.shape[-2]
        self._t = mean.shape[-1]

    def covar(self):
        raise NotImplementedError("Explicit construction of covar not implemented")

    def mean(self):
        return self._mean

    def representation(self):
        return self._mean, self._covar

    def sample(self, n_samples):
        samples = (
            self
            ._covar
            .zero_mean_mvn_samples(n_samples)
            .view(self._t, self._n, n_samples)
            .transpose(0, 1)
            .contiguous()
            .add(self._mean.unsqueeze(-1))
        )
        return samples

    def var(self):
        return self._covar.diag().view(self._n, self._t)

[jacobrgardner]

Maybe labels should always be n x 1 and then b x n x t means b batches of n data on t tasks?