Repository: cornellius-gp/gpytorch
Pull request #1029 (2020-01-23 16:56:24)

Bugfixes: lazy kernel with batch dimensions indexed correctly, error message
Fixed two small bugs I found today.

First, a `ValueError(None)` used to be thrown because of an incorrect `if` statement.

Then, a failure to replace `row`->`col` everywhere led to wrong behaviour for indexing kernel matrices. Test is also attached.


Patch:
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
index 3e322ae1..d80da585 100644
--- a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
@@ -100,14 +100,14 @@ class LazyEvaluatedKernelTensor(LazyTensor):
                 x1 = x1.expand(*([1] * len(batch_indices)), *self.x1.shape[-2:])
                 x1 = x1[(*batch_indices, row_index, dim_index)]
 
-        # Call x2[*batch_indices, row_index, :]
+        # Call x2[*batch_indices, col_index, :]
         try:
             x2 = x2[(*batch_indices, col_index, dim_index)]
         # We're going to handle multi-batch indexing with a try-catch loop
         # This way - in the default case, we can avoid doing expansions of x1 which can be timely
         except IndexError:
             if isinstance(batch_indices, slice):
-                x2 = x2.expand(1, *self.x2.shape[-2:])[(*batch_indices, row_index, dim_index)]
+                x2 = x2.expand(1, *self.x2.shape[-2:])[(*batch_indices, col_index, dim_index)]
             elif isinstance(batch_indices, tuple):
                 if any([not isinstance(bi, slice) for bi in batch_indices]):
                     raise RuntimeError(
@@ -115,7 +115,7 @@ class LazyEvaluatedKernelTensor(LazyTensor):
                         "Got batch index {batch_indices} but my shape was {self.shape}"
                     )
                 x2 = x2.expand(*([1] * len(batch_indices)), *self.x2.shape[-2:])
-                x2 = x2[(*batch_indices, row_index, dim_index)]
+                x2 = x2[(*batch_indices, col_index, dim_index)]
 
         if len(batch_indices) == 0 or all(ind == slice(None, None, None) for ind in batch_indices):
             new_kernel = self.kernel  # Avoid unnecessary copying when we aren't explicitly indexing batch dims
diff --git a/gpytorch/utils/broadcasting.py b/gpytorch/utils/broadcasting.py
index 5e11ae2a..0d650e92 100644
--- a/gpytorch/utils/broadcasting.py
+++ b/gpytorch/utils/broadcasting.py
@@ -41,7 +41,7 @@ def _matmul_broadcast_shape(shape_a, shape_b, error_msg=None):
         return shape_a[:-1]
 
     if n != shape_b[-2]:
-        if error_msg is not None:
+        if error_msg is None:
             raise RuntimeError(f"Incompatible dimensions for matmul: {shape_a} and {shape_b}")
         else:
             raise RuntimeError(error_msg)


Test patch:
diff --git a/test/lazy/test_lazy_evaluated_kernel_tensor.py b/test/lazy/test_lazy_evaluated_kernel_tensor.py
index a0de752f..abd80e78 100644
--- a/test/lazy/test_lazy_evaluated_kernel_tensor.py
+++ b/test/lazy/test_lazy_evaluated_kernel_tensor.py
@@ -42,6 +42,16 @@ class TestLazyEvaluatedKernelTensorBatch(LazyTensorTestCase, unittest.TestCase):
                     ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
                 )
 
+    def test_batch_getitem(self):
+        """Indexing was wrong when the kernel had more batch dimensions than the
+        data"""
+        x1 = torch.randn(5, 3)
+        x2 = torch.randn(5, 3)
+        kern = gpytorch.kernels.RBFKernel(batch_shape=torch.Size([2]))
+        k = kern(x1, x2)
+        self.assertEqual(k.size(), torch.Size([2, 5, 5]))
+        self.assertEqual(k[..., :4, :3].size(), torch.Size([2, 4, 3]))
+
     def test_getitem_tensor_index(self):
         # Not supported a.t.m. with LazyEvaluatedKernelTensors
         pass