Repository: cornellius-gp/gpytorch

Issue:
[Bug]
# 🐛 Bug: Possible error with multitask learning with additive kernel structure

<!-- A clear and concise description of what the bug is. -->
When I define in the class MultitaskGPModel the multitask kernel

            self.covar_module = (gpytorch.kernels.ScaleKernel(
                     gpytorch.kernels.PeriodicKernel(batch_shape=torch.Size([num_latents])) * 
                     gpytorch.kernels.RQKernel(batch_shape=torch.Size([num_latents])),
                                               batch_shape=torch.Size([num_latents])
                                                             )  + 
                                gpytorch.kernels.ScaleKernel(
                     gpytorch.kernels.MaternKernel(nu=0.5, batch_shape=torch.Size([num_latents])), 
                                               batch_shape=torch.Size([num_latents])
                                                             )
                                )

which uses the additive kernel as its outermost layer, and I apply the class on data as

  w_l = 50
  num_latents = 24
  Xc_t_npa = np.arange(0,w_l,1,dtype=np.float32).reshape(-1, 1)
  Xc_t = torch.from_numpy(Xc_t_npa).type(torch.Tensor)
  model_mul(Xc_t)

I get 
'RuntimeError: The expected shape of the kernel was torch.Size([100, 100]), but got torch.Size([24, 100, 100]). This is likely a bug in GPyTorch'.
This behavior seems not to change when changing the number of tasks or the number of latent gps.

If I use the same kernel in a non-batch setting, it works smoothly.

I wrote the batched problem with another kernel which is mathematically the same but which doesn't use the outer additive kernel, and it works smoothly. Unfortunatly the role of the subkernel parameters in the new form is not the same as that of the malfunctioning kernel, and I have to re-run a lot of past non-batch fits in the new form to make them comparable with the new setting.


## To reproduce

** Code snippet to reproduce **
```python
# Your code goes here
# Please make sure it does not require any external dependencies (other than PyTorch!)
# (We much prefer small snippets rather than links to existing libraries!)
```
    Zc_intra_np = np.arange(0, 24, 1).reshape(-1, 1)
    Zc_intra = torch.tensor(Zc_intra_np, dtype=torch.float)
 
    w_l = 50
    num_latents = 24
    num_tasks = 12
    Xc_t_npa = np.arange(0,w_l,1,dtype=np.float32).reshape(-1, 1)
    Xc_t = torch.from_numpy(Xc_t_npa).type(torch.Tensor)

    model_mul = MultitaskGPModel()
    likelihood_mul = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)
    model_mul(Xc_t)


    class MultitaskGPModel(gpytorch.models.ApproximateGP):
        
        def __init__(self):
            

            inducing_points = Zc_intra
        
            variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
                inducing_points.size(-2), batch_shape=torch.Size([num_latents])
            )

            variational_strategy = gpytorch.variational.LMCVariationalStrategy(
                gpytorch.variational.VariationalStrategy(
                    self, inducing_points, variational_distribution, learn_inducing_locations=True
                ),
                num_tasks=num_tasks,
                num_latents=num_latents,
                # could be 0
                latent_dim=-1
            )

            super().__init__(variational_strategy)

            self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([num_latents]))
        
            self.covar_module = gpytorch.kernels.ScaleKernel(
                     gpytorch.kernels.PeriodicKernel(batch_shape=torch.Size([num_latents])) * 
                     gpytorch.kernels.RQKernel(batch_shape=torch.Size([num_latents]))  + 
                     gpytorch.kernels.ScaleKernel(
                         gpytorch.kernels.MaternKernel(nu=0.5, batch_shape=torch.Size([num_latents])), 
                                                 batch_shape=torch.Size([num_latents])),                 
                             batch_shape=torch.Size([num_latents])  
                                                             )


        def forward(self, x):
            mean_x = self.mean_module(x)
            covar_x = self.covar_module(x)
            return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)






** Stack trace/error message **
```
Traceback (most recent call last):

  File "<ipython-input-398-5fc832e3a3f0>", line 1, in <module>
    model_mul(Xc_t)

  File "C:\Users\lucheron\Anaconda3\envs\pyro16_py37\lib\site-packages\gpytorch\models\approximate_gp.py", line 81, in __call__
    return self.variational_strategy(inputs, prior=prior, **kwargs)

  File "C:\Users\lucheron\Anaconda3\envs\pyro16_py37\lib\site-packages\gpytorch\variational\lmc_variational_strategy.py", line 124, in __call__
    function_dist = self.base_variational_strategy(x, prior=prior, **kwargs)

  File "C:\Users\lucheron\Anaconda3\envs\pyro16_py37\lib\site-packages\gpytorch\variational\variational_strategy.py", line 168, in __call__
    return super().__call__(x, prior=prior, **kwargs)

  File "C:\Users\lucheron\Anaconda3\envs\pyro16_py37\lib\site-packages\gpytorch\variational\_variational_strategy.py", line 129, in __call__
    **kwargs,

  File "C:\Users\lucheron\Anaconda3\envs\pyro16_py37\lib\site-packages\gpytorch\module.py", line 28, in __call__
    outputs = self.forward(*inputs, **kwargs)

  File "C:\Users\lucheron\Anaconda3\envs\pyro16_py37\lib\site-packages\gpytorch\variational\variational_strategy.py", line 96, in forward
    induc_induc_covar = full_covar[..., :num_induc, :num_induc].add_jitter()

  File "C:\Users\lucheron\Anaconda3\envs\pyro16_py37\lib\site-packages\gpytorch\lazy\lazy_evaluated_kernel_tensor.py", line 237, in add_jitter
    return self.evaluate_kernel().add_jitter(jitter_val)

  File "C:\Users\lucheron\Anaconda3\envs\pyro16_py37\lib\site-packages\gpytorch\utils\memoize.py", line 59, in g
    return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)

  File "C:\Users\lucheron\Anaconda3\envs\pyro16_py37\lib\site-packages\gpytorch\lazy\lazy_evaluated_kernel_tensor.py", line 291, in evaluate_kernel
    f"The expected shape of the kernel was {self.shape}, but got {res.shape}. "

RuntimeError: The expected shape of the kernel was torch.Size([100, 100]), but got torch.Size([24, 100, 100]). This is likely a bug in GPyTorch.
```

## Expected Behavior

<!-- A clear and concise description of what you expected to happen. --> Run with no errors

## System information

**Please complete the following information:**
- <!-- GPyTorch Version (run `print(gpytorch.__version__)` --> 1.4.1
- <!-- PyTorch Version (run `print(torch.__version__)` --> 1.8.1
- <!-- Computer OS --> Win10 pro 19042.1052

## Additional context
Add any other context about the problem here.



Discussion:
I cannot reproduce your error based on the code example you provide. I had to fix the code example to get it to run:

```python
import gpytorch
import numpy as np
import torch


class MultitaskGPModel(gpytorch.models.ApproximateGP):
    
    def __init__(self):
        

        inducing_points = Zc_intra
    
        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
            inducing_points.size(-2), batch_shape=torch.Size([num_latents])
        )

        variational_strategy = gpytorch.variational.LMCVariationalStrategy(
            gpytorch.variational.VariationalStrategy(
                self, inducing_points, variational_distribution, learn_inducing_locations=True
            ),
            num_tasks=num_tasks,
            num_latents=num_latents,
            # could be 0
            latent_dim=-1
        )

        super().__init__(variational_strategy)

        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([num_latents]))
    
        self.covar_module = gpytorch.kernels.ScaleKernel(
                 gpytorch.kernels.PeriodicKernel(batch_shape=torch.Size([num_latents])) * 
                 gpytorch.kernels.RQKernel(batch_shape=torch.Size([num_latents]))  + 
                 gpytorch.kernels.ScaleKernel(
                     gpytorch.kernels.MaternKernel(nu=0.5, batch_shape=torch.Size([num_latents])), 
                                             batch_shape=torch.Size([num_latents])),                 
                         batch_shape=torch.Size([num_latents])  
                                                         )


    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
    
    

Zc_intra_np = np.arange(0, 24, 1).reshape(-1, 1)
Zc_intra = torch.tensor(Zc_intra_np, dtype=torch.float)

w_l = 50
num_latents = 24
num_tasks = 12
Xc_t_npa = np.arange(0,w_l,1,dtype=np.float32).reshape(-1, 1)
Xc_t = torch.from_numpy(Xc_t_npa).type(torch.Tensor)

model_mul = MultitaskGPModel()
likelihood_mul = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)
model_mul(Xc_t).covariance_matrix
```
My mistake. I posted the fixed up version of the kernel - the one in which I eliminated the external additive layer. 

I gather hereafter a more elaborated, self-contained extract of the program which gives me the error.

One can run it with the 'fixed_kernel = False' flag (inside the MultitaskGPModel class) and get the error, or with the 'fixed_kernel = True' flag, not getting the error.

Thank you for the quick answer, and in perspective for your patience :)

```
import gpytorch
import numpy as np
import torch

class MultitaskGPModel(gpytorch.models.ApproximateGP):
    
    def __init__(self):
        
        inducing_points = Zc_intra
    
        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
            inducing_points.size(-2), batch_shape=torch.Size([num_latents])
        )

        variational_strategy = gpytorch.variational.LMCVariationalStrategy(
            gpytorch.variational.VariationalStrategy(
                self, inducing_points, variational_distribution, learn_inducing_locations=True
            ),
            num_tasks=num_tasks,
            num_latents=num_latents,
            latent_dim=-1
        )

        super().__init__(variational_strategy)

        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([num_latents]))
 
        
        fixed_kernel = False
    
        if fixed_kernel:            
            # fixed kernel - works
            self.covar_module = gpytorch.kernels.ScaleKernel(
                     gpytorch.kernels.PeriodicKernel(batch_shape=torch.Size([num_latents])) * 
                     gpytorch.kernels.RQKernel(batch_shape=torch.Size([num_latents]))  + 
                     gpytorch.kernels.ScaleKernel(
                         gpytorch.kernels.MaternKernel(nu=0.5, batch_shape=torch.Size([num_latents])), 
                                                 batch_shape=torch.Size([num_latents])),                 
                             batch_shape=torch.Size([num_latents])  
                                                             )
                        
        else:
            # original kernel - gives problem
            self.covar_module = gpytorch.kernels.ScaleKernel(
                     gpytorch.kernels.PeriodicKernel(batch_shape=torch.Size([num_latents])) * 
                     gpytorch.kernels.RQKernel(batch_shape=torch.Size([num_latents])),
                                               batch_shape=torch.Size([num_latents])
                                                             )  + \
                            gpytorch.kernels.ScaleKernel(
                     gpytorch.kernels.MaternKernel(nu=0.5, batch_shape=torch.Size([num_latents])), 
                                               batch_shape=torch.Size([num_latents])
                                                             )
        
                                    
    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
        
Zc_intra_np = np.arange(0, 24, 1).reshape(-1, 1)
Zc_intra = torch.tensor(Zc_intra_np, dtype=torch.float)

w_l = 50
num_latents = 24
num_tasks = 12
Xc_t_npa = np.arange(0,w_l,1,dtype=np.float32).reshape(-1, 1)
Xc_t = torch.from_numpy(Xc_t_npa).type(torch.Tensor)

model_mul = MultitaskGPModel()
likelihood_mul = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)

model_mul(Xc_t).covariance_matrix

Yc_t = torch.rand((50,12))
    
Yc_t = torch.tensor(Yc_t,dtype=torch.float
                    ).flatten().reshape(Yc_t.shape[0],Yc_t.shape[1])

training_iter = 5

model_mul.train()
likelihood_mul.train()

optimizer_mul = torch.optim.Adam([
    {'params': model_mul.parameters()},
    {'params': likelihood_mul.parameters()},
], lr=0.1)

mll_mul = gpytorch.mlls.VariationalELBO(likelihood_mul, model_mul, num_data=Yc_t.size(0))

for i in range(training_iter):    
        optimizer_mul.zero_grad()
        output = model_mul(Xc_t)
        loss = -mll_mul(output, Yc_t)
        loss.backward()
        optimizer_mul.step()
```

The error is now
```
RuntimeError: The expected shape of the kernel was torch.Size([24, 24]), but got torch.Size([24, 24, 24]). This is likely a bug in GPyTorch.
```
and can be obtained by just running
```
model_mul(Xc_t)
```
after the model is instantiated.

One extra piece of information, which I don't know if it can be useful. In my environment I first installed pyro-ppl 1.6.0, then gpytorch 1.4.1.


Further information.
I built from scratch a new Python/Spyder environment with GPyTorch 1.4.1 without previous Pyro installation. 
Running the program with the 'fixed_kernel = False' returns again the 'This is likely a bug in GPyTorch.' message. 
Program runs smoothly wiith  'fixed_kernel = True'.



Resolving patch:
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -331,8 +331,8 @@ def covar_dist(
         return res
 
     def named_sub_kernels(self):
-        for name, module in self._modules.items():
-            if isinstance(module, Kernel):
+        for name, module in self.named_modules():
+            if module is not self and isinstance(module, Kernel):
                 yield name, module
 
     def num_outputs_per_input(self, x1, x2):


Test patch:
diff --git a/test/kernels/test_additive_and_product_kernels.py b/test/kernels/test_additive_and_product_kernels.py
--- a/test/kernels/test_additive_and_product_kernels.py
+++ b/test/kernels/test_additive_and_product_kernels.py
@@ -56,6 +56,23 @@ def test_computes_product_of_radial_basis_function(self):
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 2e-5)
 
+    def test_computes_product_of_radial_basis_function_batch(self):
+        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
+        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
+        lengthscale = 2
+
+        kernel_1 = RBFKernel(batch_shape=torch.Size([4])).initialize(lengthscale=lengthscale)
+        kernel_2 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel = kernel_1 * kernel_2
+
+        actual = torch.tensor([[16, 4], [4, 0], [64, 36]], dtype=torch.float)
+        actual = actual.mul_(-0.5).div_(lengthscale ** 2).exp() ** 2
+        actual = actual.repeat(4, 1, 1)
+
+        kernel.eval()
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 2e-5)
+
     def test_computes_sum_of_radial_basis_function(self):
         a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
         b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)