Repository: cornellius-gp/gpytorch (python)

Issue:
batch GridInterpolationKernel [Feature Request]
# 🚀 Feature Request

Add support for batched SKI kernel approximations. I'm not sure how one would go about adding this functionality. Ideally `GridInterpolationKernel` would infer the batch shape from `base_kernel`.

If this can't be done soon, we should add a warning that GridInterpolationKernel doesn't support batch evaluation. Currently it seems to fail silently by producing incorrectly shaped outputs. 

```
>>> import torch
>>> import gpytorch
>>> base_kernel = gpytorch.kernels.RBFKernel(batch_shape=torch.Size([2]))
>>> ski_kernel = gpytorch.kernels.GridInterpolationKernel(base_kernel, num_dims=1, grid_size=128)
>>> a = torch.rand(2)
>>> ski_kernel(a, a).shape
torch.Size([2, 2])
>>> base_kernel(a, a).shape
torch.Size([2, 2, 2])
```

## Motivation

In general for multi-dimensional outputs we have to train and evaluate multiple GPs in parallel. Currently this functionality appears to be completely absent in `GridInterpolationKernel`. I think SKI could have compelling applications in the context of online model learning, where one may be unable or unwilling to solve a slowly converging variational optimization problem.

## Pitch

Ultimately it would be amazing if `batch_shape` could be added as an argument to the `GridInterpolationKernel` constructor, or inferred from the base kernel. I'm not sure how much work this would involve.

One alternative could be relying on a multitask-like approach, but this requires memory usage that scales quadratically with the output dimension, which quickly becomes impractical. 

I'd be willing to open a pull request, but I'd need some pointers for how to get started. How was arbitrary batch_shape supported in the first place?


Pull request description:
Fix grid kernel batch issues.
This addresses using different batch shapes within GridInterpolationKernel. Now it's possible to use any sort of broadcastable batch sizes with  1) input data, 2) inducing points, and 3) kernel hypers.

[Fixes #963]

Patch:
diff --git a/gpytorch/kernels/grid_interpolation_kernel.py b/gpytorch/kernels/grid_interpolation_kernel.py
index d7a9f3f2..48c2723a 100644
--- a/gpytorch/kernels/grid_interpolation_kernel.py
+++ b/gpytorch/kernels/grid_interpolation_kernel.py
@@ -6,6 +6,7 @@ import torch
 
 from ..lazy import InterpolatedLazyTensor
 from ..models.exact_prediction_strategies import InterpolatedPredictionStrategy
+from ..utils.broadcasting import _mul_broadcast_shape
 from ..utils.grid import create_grid
 from ..utils.interpolation import Interpolation
 from .grid_kernel import GridKernel
@@ -174,8 +175,6 @@ class GridInterpolationKernel(GridKernel):
         base_lazy_tsr = self._inducing_forward(last_dim_is_batch=last_dim_is_batch, **params)
         if last_dim_is_batch:
             base_lazy_tsr = base_lazy_tsr.repeat(*x1.shape[:-2], x1.size(-1), 1, 1)
-        if x1.dim() > 2:
-            base_lazy_tsr = base_lazy_tsr.repeat(*x1.shape[:-2], 1, 1)
 
         left_interp_indices, left_interp_values = self._compute_grid(x1, last_dim_is_batch)
         if torch.equal(x1, x2):
@@ -184,12 +183,15 @@ class GridInterpolationKernel(GridKernel):
         else:
             right_interp_indices, right_interp_values = self._compute_grid(x2, last_dim_is_batch)
 
+        batch_shape = _mul_broadcast_shape(
+            base_lazy_tsr.batch_shape, left_interp_indices.shape[:-2], right_interp_indices.shape[:-2]
+        )
         res = InterpolatedLazyTensor(
-            base_lazy_tsr,
-            left_interp_indices.detach(),
-            left_interp_values,
-            right_interp_indices.detach(),
-            right_interp_values,
+            base_lazy_tsr.expand(*batch_shape, *base_lazy_tsr.matrix_shape),
+            left_interp_indices.detach().expand(*batch_shape, *left_interp_indices.shape[-2:]),
+            left_interp_values.expand(*batch_shape, *left_interp_values.shape[-2:]),
+            right_interp_indices.detach().expand(*batch_shape, *right_interp_indices.shape[-2:]),
+            right_interp_values.expand(*batch_shape, *right_interp_values.shape[-2:]),
         )
 
         if diag:
diff --git a/gpytorch/kernels/grid_kernel.py b/gpytorch/kernels/grid_kernel.py
index d725642f..da55dd18 100644
--- a/gpytorch/kernels/grid_kernel.py
+++ b/gpytorch/kernels/grid_kernel.py
@@ -51,7 +51,7 @@ class GridKernel(Kernel):
         if not base_kernel.is_stationary:
             raise RuntimeError("The base_kernel for GridKernel must be stationary.")
 
-        super(GridKernel, self).__init__(active_dims=active_dims)
+        super().__init__(active_dims=active_dims)
         if torch.is_tensor(grid):
             grid = convert_legacy_grid(grid)
         self.interpolation_mode = interpolation_mode
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index 50cd12a0..30d4662f 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -12,6 +12,7 @@ from ..constraints import Positive
 from ..lazy import LazyEvaluatedKernelTensor, ZeroLazyTensor, delazify, lazify
 from ..models.exact_prediction_strategies import DefaultPredictionStrategy, SumPredictionStrategy
 from ..module import Module
+from ..utils.broadcasting import _mul_broadcast_shape
 from ..utils.deprecation import _ClassWithDeprecatedBatchSize, _deprecate_kwarg_with_transform
 
 
@@ -129,13 +130,6 @@ class Kernel(Module, _ClassWithDeprecatedBatchSize):
 
     has_lengthscale = False
 
-    @property
-    def is_stationary(self) -> bool:
-        """
-        Property to indicate whether kernel is stationary or not.
-        """
-        return self.has_lengthscale
-
     def __init__(
         self,
         ard_num_dims=None,
@@ -154,7 +148,7 @@ class Kernel(Module, _ClassWithDeprecatedBatchSize):
         self.register_buffer("active_dims", active_dims)
         self.ard_num_dims = ard_num_dims
 
-        self.batch_shape = _deprecate_kwarg_with_transform(
+        self._batch_shape = _deprecate_kwarg_with_transform(
             kwargs, "batch_size", "batch_shape", batch_shape, lambda n: torch.Size([n])
         )
 
@@ -216,6 +210,18 @@ class Kernel(Module, _ClassWithDeprecatedBatchSize):
         """
         raise NotImplementedError()
 
+    @property
+    def batch_shape(self):
+        kernels = list(self.sub_kernels())
+        if len(kernels):
+            return _mul_broadcast_shape(self._batch_shape, *[k.batch_shape for k in kernels])
+        else:
+            return self._batch_shape
+
+    @batch_shape.setter
+    def batch_shape(self, val):
+        self._batch_shape = val
+
     @property
     def dtype(self):
         if self.has_lengthscale:
@@ -225,6 +231,13 @@ class Kernel(Module, _ClassWithDeprecatedBatchSize):
                 return param.dtype
             return torch.get_default_dtype()
 
+    @property
+    def is_stationary(self) -> bool:
+        """
+        Property to indicate whether kernel is stationary or not.
+        """
+        return self.has_lengthscale
+
     @property
     def lengthscale(self):
         if self.has_lengthscale:
@@ -322,6 +335,11 @@ class Kernel(Module, _ClassWithDeprecatedBatchSize):
 
         return res
 
+    def named_sub_kernels(self):
+        for name, module in self._modules.items():
+            if isinstance(module, Kernel):
+                yield name, module
+
     def num_outputs_per_input(self, x1, x2):
         """
         How many outputs are produced per input (default 1)
@@ -334,6 +352,10 @@ class Kernel(Module, _ClassWithDeprecatedBatchSize):
     def prediction_strategy(self, train_inputs, train_prior_dist, train_labels, likelihood):
         return DefaultPredictionStrategy(train_inputs, train_prior_dist, train_labels, likelihood)
 
+    def sub_kernels(self):
+        for _, kernel in self.named_sub_kernels():
+            yield kernel
+
     def __call__(self, x1, x2=None, diag=False, last_dim_is_batch=False, **params):
         x1_, x2_ = x1, x2
 
@@ -407,9 +429,8 @@ class Kernel(Module, _ClassWithDeprecatedBatchSize):
             new_batch_shape_len = len(self.batch_shape) - ndim_removed
             new_kernel.batch_shape = new_kernel._parameters[param_name].shape[:new_batch_shape_len]
 
-        for sub_module_name, sub_module in self._modules.items():
-            if isinstance(sub_module, Kernel):
-                self._modules[sub_module_name] = sub_module.__getitem__(index)
+        for sub_module_name, sub_module in self.named_sub_kernels():
+            self._modules[sub_module_name] = sub_module.__getitem__(index)
 
         return new_kernel
 


Test patch:
diff --git a/test/kernels/test_grid_interpolation_kernel.py b/test/kernels/test_grid_interpolation_kernel.py
new file mode 100644
index 00000000..94987a04
--- /dev/null
+++ b/test/kernels/test_grid_interpolation_kernel.py
@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+
+import unittest
+
+import torch
+
+from gpytorch.kernels import GridInterpolationKernel, RBFKernel
+from gpytorch.lazy import InterpolatedLazyTensor
+
+
+class TestGridInterpolationKernel(unittest.TestCase):
+    def test_standard(self):
+        base_kernel = RBFKernel()
+        kernel = GridInterpolationKernel(base_kernel, num_dims=2, grid_size=128, grid_bounds=[(-1.2, 1.2)] * 2)
+
+        xs = torch.randn(5, 2).clamp(-1, 1)
+        interp_covar = kernel(xs, xs).evaluate_kernel()
+        self.assertIsInstance(interp_covar, InterpolatedLazyTensor)
+
+        xs = torch.randn(5, 2).clamp(-1, 1)
+        grid_eval = kernel(xs, xs).evaluate()
+        actual_eval = base_kernel(xs, xs).evaluate()
+        self.assertLess(torch.norm(grid_eval - actual_eval), 2e-5)
+
+        xs = torch.randn(3, 5, 2).clamp(-1, 1)
+        grid_eval = kernel(xs, xs).evaluate()
+        actual_eval = base_kernel(xs, xs).evaluate()
+        self.assertLess(torch.norm(grid_eval - actual_eval), 2e-5)
+
+    def test_batch_base_kernel(self):
+        base_kernel = RBFKernel(batch_shape=torch.Size([3]))
+        kernel = GridInterpolationKernel(base_kernel, num_dims=2, grid_size=128, grid_bounds=[(-1.2, 1.2)] * 2)
+
+        xs = torch.randn(5, 2).clamp(-1, 1)
+        grid_eval = kernel(xs, xs).evaluate()
+        actual_eval = base_kernel(xs, xs).evaluate()
+        self.assertLess(torch.norm(grid_eval - actual_eval), 2e-5)
+
+        xs = torch.randn(3, 5, 2).clamp(-1, 1)
+        grid_eval = kernel(xs, xs).evaluate()
+        actual_eval = base_kernel(xs, xs).evaluate()
+        self.assertLess(torch.norm(grid_eval - actual_eval), 2e-5)
+
+        xs = torch.randn(4, 3, 5, 2).clamp(-1, 1)
+        grid_eval = kernel(xs, xs).evaluate()
+        actual_eval = base_kernel(xs, xs).evaluate()
+        self.assertLess(torch.norm(grid_eval - actual_eval), 2e-5)
+
+
+if __name__ == "__main__":
+    unittest.main()