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.


Changed files:
gpytorch/kernels/kernel.py, test/kernels/test_additive_and_product_kernels.py

diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index 3868597d..7c70bbfb 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -331,8 +331,8 @@ class Kernel(Module):
         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):
diff --git a/test/kernels/test_additive_and_product_kernels.py b/test/kernels/test_additive_and_product_kernels.py
index e2d3b0d5..a7c52747 100644
--- a/test/kernels/test_additive_and_product_kernels.py
+++ b/test/kernels/test_additive_and_product_kernels.py
@@ -56,6 +56,23 @@ class TestAdditiveAndProductKernel(unittest.TestCase):
         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)