Repository: cornellius-gp/gpytorch

Issue:
Problem with batch Hadamard Multitask GP regression (likely a bug in GPyTorch)
Hi I was trying to build a batch hadamard multitask GP regression model following this:
https://gpytorch.readthedocs.io/en/latest/examples/03_Multitask_Exact_GPs/Hadamard_Multitask_GP_Regression.html 

Here is my code: 

    import math
    import torch
    import gpytorch
    from matplotlib import pyplot as plt

    # %%

    train_x1 = torch.rand(50)
    train_x2 = torch.rand(10)

    train_y1 = torch.sin(train_x1 * (2 * math.pi)) + \
        torch.randn(train_x1.size()) * 0.2
    train_y2 = torch.cos(train_x2 * (2 * math.pi)) + \
        torch.randn(train_x2.size()) * 0.2

    # %%


    class MultitaskGPModel(gpytorch.models.ExactGP):
        def __init__(self, train_x, train_y, likelihood):
            super(MultitaskGPModel, self).__init__(train_x, train_y, likelihood)
            self.mean_module = gpytorch.means.ConstantMean(
                batch_shape=torch.Size([4]))
            self.covar_module = gpytorch.kernels.RBFKernel(
                batch_shape=torch.Size([4]))

            self.task_covar_module = gpytorch.kernels.IndexKernel(
                num_tasks=2, rank=1, batch_shape=torch.Size([4]))

        def forward(self, x, i):
            mean_x = self.mean_module(x)

            # Get input-input covariance
            covar_x = self.covar_module(x)
            # Get task-task covariance
            covar_i = self.task_covar_module(i)
            # Multiply the two together to get the covariance we want
            covar = covar_x.mul(covar_i)

            return gpytorch.distributions.MultivariateNormal(mean_x, covar)


    likelihood = gpytorch.likelihoods.GaussianLikelihood()

    # %% Reshaping to simulate batches
    train_i_task1 = torch.full_like(train_x1, dtype=torch.long, fill_value=0)
    train_i_task2 = torch.full_like(train_x2, dtype=torch.long, fill_value=1)

    full_train_x = torch.cat([train_x1, train_x2])
    full_train_x = full_train_x.view(1, -1, 1).repeat(4, 1, 1)

    full_train_i = torch.cat([train_i_task1, train_i_task2])
    full_train_i = full_train_i.view(1, -1, 1).repeat(4, 1, 1)

    full_train_y = torch.cat([train_y1, train_y2])
    full_train_y = full_train_y.unsqueeze(0)
    full_train_y = torch.cat(
        (full_train_y, full_train_y, full_train_y, full_train_y))
        
    # Here we have two items that we're passing in as train_inputs
    model = MultitaskGPModel((full_train_x, full_train_i),
                             full_train_y, likelihood)

    # %%

    training_iterations = 50

    # Find optimal model hyperparameters
    model.train()
    likelihood.train()

    # Use the adam optimizer
    optimizer = torch.optim.Adam([
        {'params': model.parameters()},  # Includes GaussianLikelihood parameters
    ], lr=0.1)

    # "Loss" for GPs - the marginal log likelihood
    mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

    for i in range(training_iterations):
        optimizer.zero_grad()
        output = model(full_train_x, full_train_i)
        loss = -mll(output, full_train_y).sum()
        loss.backward()
        print('Iter %d/50 - Loss: %.3f' % (i + 1, loss.item()))
        optimizer.step()

I just replicated the existing train x, i and y values 4 times to get a batch size of 4. 

However when I run this, it spits out an error: 

    ---------------------------------------------------------------------------
    RuntimeError                              Traceback (most recent call last)
    ~/test_multitask_gp_unequal_samples.py in 
         87 for i in range(training_iterations):
         88     optimizer.zero_grad()
    ---> 89     output = model(full_train_x, full_train_i)
         90     loss = -mll(output, full_train_y).sum()
         91     loss.backward()

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
        263                 if not all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):
        264                     raise RuntimeError("You must train on the training inputs!")
    --> 265             res = super().__call__(*inputs, **kwargs)
        266             return res
        267 

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)
         22 
         23     def __call__(self, *inputs, **kwargs):
    ---> 24         outputs = self.forward(*inputs, **kwargs)
         25         if isinstance(outputs, list):
         26             return [_validate_module_outputs(output) for output in outputs]

    ~/forecast_research/general_experiments/probabilistic_programs/test_multitask_gp_unequal_samples.py in forward(self, x, i)
         44         covar_i = self.task_covar_module(i)
         45         # Multiply the two together to get the covariance we want
    ---> 46         covar = covar_x.mul(covar_i)
         47 
         48         return gpytorch.distributions.MultivariateNormal(mean_x, covar)

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in mul(self, other)
       1145                 return self._mul_constant(other.view(*other.shape[:-2]))
       1146 
    -> 1147         return self._mul_matrix(lazify(other))
       1148 
       1149     def ndimension(self):

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in _mul_matrix(self, other)
        502 
        503         self = self.evaluate_kernel()
    --> 504         other = other.evaluate_kernel()
        505         if isinstance(self, NonLazyTensor) or isinstance(other, NonLazyTensor):
        506             return NonLazyTensor(self.evaluate() * other.evaluate())

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
         32         cache_name = name if name is not None else method
         33         if not is_in_cache(self, cache_name):
    ---> 34             add_to_cache(self, cache_name, method(self, *args, **kwargs))
         35         return get_from_cache(self, cache_name)
         36 

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate_kernel(self)
        283             if res.shape != self.shape:
        284                 raise RuntimeError(
    --> 285                     f"The expected shape of the kernel was {self.shape}, but got {res.shape}. "
        286                     "This is likely a bug in GPyTorch."
        287                 )

    RuntimeError: The expected shape of the kernel was torch.Size([4, 60, 60]), but got torch.Size([4, 4, 60, 60]). This is likely a bug in GPyTorch. 

I think there is problem when it tries to multiply the two covariance matrices. 

Please let me know what the possible solution could be for this. 

I am using GPyTorch v1.0.0 and PyTorch v1.4


Thanks 



Discussion:
Sorry for the delay looking at this @grishabhg -- things are a bit crazy!

This does indeed appear to likely be a bug on our end. I will investigate and see if I can't put up a PR in the near future.
Great! thanks @jacobrgardner 
Please do let me know when you guys solve it. 

Resolving patch:
diff --git a/gpytorch/kernels/index_kernel.py b/gpytorch/kernels/index_kernel.py
--- a/gpytorch/kernels/index_kernel.py
+++ b/gpytorch/kernels/index_kernel.py
@@ -80,9 +80,10 @@ def covar_matrix(self):
 
     def forward(self, i1, i2, **params):
         covar_matrix = self._eval_covar_matrix()
+
         res = InterpolatedLazyTensor(
             base_lazy_tensor=covar_matrix,
-            left_interp_indices=i1.expand(self.batch_shape + i1.shape),
-            right_interp_indices=i2.expand(self.batch_shape + i2.shape),
+            left_interp_indices=i1.expand(self.batch_shape + i1.shape[-2:]),
+            right_interp_indices=i2.expand(self.batch_shape + i2.shape[-2:]),
         )
         return res


Test patch: