Fix LazyEvaluatedKernelTensor._unsqueeze_batch (#1828)

Different implementation of `LazyEvaluatedKernelTensor._unsqueeze_batch`

Co-authored-by: Geoff Pleiss <[email protected]>

Author: valtron

gpytorch/lazy/lazy_evaluated_kernel_tensor.py

@@ -68,8 +68,16 @@ def _getitem(self, row_index, col_index, *batch_indices):
             # Now we know that x1 and x2 are slices
             # Let's make sure that the slice dimensions perfectly correspond with the number of
             # outputs per input that we have
-            row_start, row_end, row_step = row_index.start, row_index.stop, row_index.step
-            col_start, col_end, col_step = col_index.start, col_index.stop, col_index.step
+            row_start, row_end, row_step = (
+                row_index.start,
+                row_index.stop,
+                row_index.step,
+            )
+            col_start, col_end, col_step = (
+                col_index.start,
+                col_index.stop,
+                col_index.step,
+            )
             if row_step is not None or col_step is not None:
                 return self.evaluate_kernel()._getitem(row_index, col_index, *batch_indices)
             if (
@@ -127,7 +135,7 @@ def _getitem(self, row_index, col_index, *batch_indices):
             new_kernel = self.kernel.__getitem__(batch_indices)
 
         # Now construct a kernel with those indices
-        return self.__class__(x1, x2, kernel=new_kernel, last_dim_is_batch=self.last_dim_is_batch, **self.params)
+        return self.__class__(x1, x2, kernel=new_kernel, last_dim_is_batch=self.last_dim_is_batch, **self.params,)
 
     def _matmul(self, rhs):
         # This _matmul is defined computes the kernel in chunks
@@ -148,7 +156,7 @@ def _matmul(self, rhs):
             res = []
             for sub_x1 in sub_x1s:
                 sub_kernel_matrix = lazify(
-                    self.kernel(sub_x1, x2, diag=False, last_dim_is_batch=self.last_dim_is_batch, **self.params)
+                    self.kernel(sub_x1, x2, diag=False, last_dim_is_batch=self.last_dim_is_batch, **self.params,)
                 )
                 res.append(sub_kernel_matrix._matmul(rhs))
 
@@ -177,7 +185,7 @@ def _quad_form_derivative(self, left_vecs, right_vecs):
             sub_x1.requires_grad_(True)
             with torch.enable_grad(), settings.lazily_evaluate_kernels(False):
                 sub_kernel_matrix = lazify(
-                    self.kernel(sub_x1, x2, diag=False, last_dim_is_batch=self.last_dim_is_batch, **self.params)
+                    self.kernel(sub_x1, x2, diag=False, last_dim_is_batch=self.last_dim_is_batch, **self.params,)
                 )
             sub_grad_outputs = tuple(sub_kernel_matrix._quad_form_derivative(sub_left_vecs, right_vecs))
             sub_kernel_outputs = tuple(sub_kernel_matrix.representation())
@@ -230,14 +238,16 @@ def _size(self):
 
     def _transpose_nonbatch(self):
         return self.__class__(
-            self.x2, self.x1, kernel=self.kernel, last_dim_is_batch=self.last_dim_is_batch, **self.params
+            self.x2, self.x1, kernel=self.kernel, last_dim_is_batch=self.last_dim_is_batch, **self.params,
         )
 
     def add_jitter(self, jitter_val=1e-3):
         return self.evaluate_kernel().add_jitter(jitter_val)
 
     def _unsqueeze_batch(self, dim):
-        return self[(slice(None),) * dim + (None,)]
+        x1 = self.x1.unsqueeze(dim)
+        x2 = self.x2.unsqueeze(dim)
+        return self.__class__(x1, x2, kernel=self.kernel, last_dim_is_batch=self.last_dim_is_batch, **self.params,)
 
     @cached(name="kernel_diag")
     def diag(self):
@@ -281,7 +291,7 @@ def evaluate_kernel(self):
         with settings.lazily_evaluate_kernels(False):
             temp_active_dims = self.kernel.active_dims
             self.kernel.active_dims = None
-            res = self.kernel(x1, x2, diag=False, last_dim_is_batch=self.last_dim_is_batch, **self.params)
+            res = self.kernel(x1, x2, diag=False, last_dim_is_batch=self.last_dim_is_batch, **self.params,)
             self.kernel.active_dims = temp_active_dims
 
         # Check the size of the output
@@ -305,7 +315,7 @@ def repeat(self, *repeats):
 
         x1 = self.x1.repeat(*batch_repeat, row_repeat, 1)
         x2 = self.x2.repeat(*batch_repeat, col_repeat, 1)
-        return self.__class__(x1, x2, kernel=self.kernel, last_dim_is_batch=self.last_dim_is_batch, **self.params)
+        return self.__class__(x1, x2, kernel=self.kernel, last_dim_is_batch=self.last_dim_is_batch, **self.params,)
 
     def representation(self):
         # If we're checkpointing the kernel, we'll use chunked _matmuls defined in LazyEvaluatedKernelTensor