Fowarding through RBFKernel in batch-mode broken

The recent changes to master have basically completely broken batch mode. This example is for the forward method in kernels.RBFKernel. As an example, consider the following code snippet

import torch
import gpytorch

class BatchGP(gpytorch.models.ExactGP):
    def __init__(self, train_inputs, train_targets, likelihood):
        super(BatchGP, self).__init__(train_inputs, train_targets, likelihood)
        self.mean_module = gpytorch.means.ConstantMean(batch_size=2)
        self.covar_module = gpytorch.kernels.RBFKernel(ard_num_dims=3, batch_size=2)
        self.register_parameter('log_outputscale', torch.nn.Parameter(torch.zeros(2, 1, 1)))
        
    def forward(self, inputs):
        inputs = inputs.unsqueeze(0)
        inputs = inputs.expand(2, -1, -1)
        mean = self.mean_module(inputs)
        prescale_covar = self.covar_module(inputs).evaluate()
        covar = self.log_outputscale.exp() * prescale_covar

        return self.likelihood(gpytorch.random_variables.GaussianRandomVariable(mean, covar))

train_x = torch.rand(5, 3)
train_y = torch.rand(5, 2)
likelihood = gpytorch.likelihoods.GaussianLikelihood()

batch_gp = BatchGP(train_x, train_y, likelihood)

batch_gp(train_x)

Out:

RuntimeError                              Traceback (most recent call last)
<ipython-input-11-bbf86dd62a62> in <module>()
     24 batch_gp = BatchGP(train_x, train_y, likelihood)
     25 
---> 26 batch_gp(train_x)

~/Code/gpytorch/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
     79                 if not all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):
     80                     raise RuntimeError("You must train on the training inputs!")
---> 81             return super(ExactGP, self).__call__(*inputs, **kwargs)
     82 
     83         # Posterior mode

~/Code/gpytorch/gpytorch/module.py in __call__(self, *inputs, **kwargs)
    178 
    179     def __call__(self, *inputs, **kwargs):
--> 180         outputs = self.forward(*inputs, **kwargs)
    181         if torch.is_tensor(outputs) or isinstance(outputs, RandomVariable) or isinstance(outputs, LazyTensor):
    182             return outputs

<ipython-input-11-bbf86dd62a62> in forward(self, inputs)
     13         inputs = inputs.expand(2, -1, -1)
     14         mean = self.mean_module(inputs)
---> 15         prescale_covar = self.covar_module(inputs).evaluate()
     16         covar = self.log_outputscale.exp() * prescale_covar
     17 

~/Code/gpytorch/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate(self)
    123 
    124     def evaluate(self):
--> 125         return self.evaluate_kernel().evaluate()
    126 
    127     def exact_predictive_mean(self, full_mean, train_labels, n_train, likelihood, precomputed_cache=None):

~/Code/gpytorch/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate_kernel(self)
    104                 x2 = self.x2
    105 
--> 106             self._cached_kernel_eval = super(Kernel, self.kernel).__call__(x1, x2, **self.params)
    107             if self.squeeze_row:
    108                 self._cached_kernel_eval.squeeze_(-2)

~/Code/gpytorch/gpytorch/module.py in __call__(self, *inputs, **kwargs)
    178 
    179     def __call__(self, *inputs, **kwargs):
--> 180         outputs = self.forward(*inputs, **kwargs)
    181         if torch.is_tensor(outputs) or isinstance(outputs, RandomVariable) or isinstance(outputs, LazyTensor):
    182             return outputs

~/Code/gpytorch/gpytorch/kernels/rbf_kernel.py in forward(self, x1, x2)
     94     def forward(self, x1, x2):
     95         x1_, x2_ = self._create_input_grid(x1, x2)
---> 96         x1_ = x1_.div(self.lengthscale)
     97         x2_ = x2_.div(self.lengthscale)
     98 

RuntimeError: The size of tensor a (5) must match the size of tensor b (2) at non-singleton dimension 1

If you go in and look at the dimension of the tensors, x1_.size() is [B x N x 1 x D] and self.lengthscale.size() is [B x 1 x D]

[samuelstanton]

@gpleiss There was indeed a mistake in that example. I’ve update the issue with an example of the issue with forwarding through kernels in batch mode. After some more investigation it seems to be something that happens specifically in covar_module(input).evaluate(). If you don’t try to convert to tensor the forward pass seems to work fine. e.g.

    def forward(self, inputs):
        inputs = inputs.unsqueeze(0)
        inputs = inputs.expand(2, -1, -1)
        mean = self.mean_module(inputs)
        covar = gpytorch.lazy.DiagLazyTensor(self.covar_module(inputs).diag())

        return self.likelihood(gpytorch.random_variables.GaussianRandomVariable(mean, covar))

But .evaluate() does work sometimes, like the following:

    def forward(self, inputs):
        inputs = inputs.unsqueeze(0)
        inputs = inputs.expand(2, -1, -1)
        mean = self.mean_module(inputs)
        prescaled_covar = gpytorch.lazy.DiagLazyTensor(self.covar_module(inputs).diag()).evaluate()
        covar = self.log_outputscale.exp() * prescaled_covar

        return self.likelihood(gpytorch.random_variables.GaussianRandomVariable(mean, covar))

So maybe a better name for the issue would be .evaluate() is unstable in batchmode?

[gpleiss]

@samuelstanton try now