Repository: cornellius-gp/gpytorch

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?



Discussion:

Resolving patch:
diff --git a/gpytorch/kernels/grid_interpolation_kernel.py b/gpytorch/kernels/grid_interpolation_kernel.py
--- a/gpytorch/kernels/grid_interpolation_kernel.py
+++ b/gpytorch/kernels/grid_interpolation_kernel.py
@@ -6,6 +6,7 @@
 
 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 @@ def forward(self, x1, x2, diag=False, last_dim_is_batch=False, **params):
         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 @@ def forward(self, x1, x2, diag=False, last_dim_is_batch=False, **params):
         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
--- a/gpytorch/kernels/grid_kernel.py
+++ b/gpytorch/kernels/grid_kernel.py
@@ -51,7 +51,7 @@ def __init__(
         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
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -12,6 +12,7 @@
 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 @@ def __init__(
         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 @@ def forward(self, x1, x2, diag=False, last_dim_is_batch=False, **params):
         """
         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 @@ def dtype(self):
                 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 @@ def covar_dist(
 
         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 @@ def num_outputs_per_input(self, x1, x2):
     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 @@ def __getitem__(self, index):
             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
--- /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()