Repository: cornellius-gp/gpytorch
Pull request #888 (2019-09-30 07:23:25)

Independent grid sizes per dimension for GridKernel and GridInterpolationKernel
This PR finishes #861 and implements independent grid sizes per dimension for GridKernel and GridInterpolationKernel. The changes are mostly the same, except that this PR makes additional changes to pass all the tests and ensure correctness. That PR had a few merges from other branches, which clutters the git log a bit, so I'm opening this new PR.

One thing that will have to be done in the future is to fully deprecate the old method of passing in a `G x D` grid into `GridKernel` where `G` is the number of grid points for every dimension. For now, we use a function that converts this old grid into the new expected format of `List[Tensor]` with each Tensor possibly having different lengths.

In addition, it's important to note that this change causes the computations for each dimension to be done sequentially instead of in parallel through batch mode. It may be useful to special case the scenario when all the grid sizes are the same in the future to ensure maximum parallelism. Still, if we assume that data dimensions for GridKernels and GridInterpolation are of reasonable sizes, say less than 10, then the for-loop slowdowns will not be that significant.

This PR also does the following
- Add previously non-existent docstrings to `ToeplitzLazyTensor`
- Swap to using `torch.meshgrid` for `gpytorch.utils.grid.create_data_from_grid` instead of for-loops
- Added type hints for `GridKernel`, `GridInterpolationKernel` and some of the related code


Patch:
diff --git a/gpytorch/kernels/grid_interpolation_kernel.py b/gpytorch/kernels/grid_interpolation_kernel.py
index 69f49c20..6377f0ef 100644
--- a/gpytorch/kernels/grid_interpolation_kernel.py
+++ b/gpytorch/kernels/grid_interpolation_kernel.py
@@ -1,10 +1,13 @@
 #!/usr/bin/env python3
 
 import torch
+from .kernel import Kernel
 from .grid_kernel import GridKernel
 from ..lazy import InterpolatedLazyTensor
 from ..models.exact_prediction_strategies import InterpolatedPredictionStrategy
 from ..utils.interpolation import Interpolation
+from typing import Union, List, Optional, Tuple
+from gpytorch.utils.grid import create_grid
 
 
 class GridInterpolationKernel(GridKernel):
@@ -47,8 +50,9 @@ class GridInterpolationKernel(GridKernel):
     Args:
         - :attr:`base_kernel` (Kernel):
             The kernel to approximate with KISS-GP
-        - :attr:`grid_size` (int):
-            The size of the grid (in each dimension)
+        - :attr:`grid_size` (Union[int, List[int]]):
+            The size of the grid in each dimension.
+            If a single int is provided, then every dimension will have the same grid size.
         - :attr:`num_dims` (int):
             The dimension of the input data. Required if `grid_bounds=None`
         - :attr:`grid_bounds` (tuple(float, float), optional):
@@ -62,7 +66,14 @@ class GridInterpolationKernel(GridKernel):
         http://proceedings.mlr.press/v37/wilson15.pdf
     """
 
-    def __init__(self, base_kernel, grid_size, num_dims=None, grid_bounds=None, active_dims=None):
+    def __init__(
+        self,
+        base_kernel: Kernel,
+        grid_size: Union[int, List[int]],
+        num_dims: int = None,
+        grid_bounds: Optional[Tuple[float, float]] = None,
+        active_dims: Tuple[int, ...] = None,
+    ):
         has_initialized_grid = 0
         grid_is_dynamic = True
 
@@ -84,31 +95,29 @@ class GridInterpolationKernel(GridKernel):
                     "grid_bounds ({})".format(num_dims, len(grid_bounds))
                 )
 
+        if isinstance(grid_size, int):
+            grid_sizes = [grid_size for _ in range(num_dims)]
+        else:
+            grid_sizes = list(grid_size)
+
+        if len(grid_sizes) != num_dims:
+            raise RuntimeError("The number of grid sizes provided through grid_size do not match num_dims.")
+
         # Initialize values and the grid
         self.grid_is_dynamic = grid_is_dynamic
         self.num_dims = num_dims
-        self.grid_size = grid_size
+        self.grid_sizes = grid_sizes
         self.grid_bounds = grid_bounds
-        grid = self._create_grid()
+        grid = create_grid(self.grid_sizes, self.grid_bounds)
 
         super(GridInterpolationKernel, self).__init__(
             base_kernel=base_kernel, grid=grid, interpolation_mode=True, active_dims=active_dims
         )
         self.register_buffer("has_initialized_grid", torch.tensor(has_initialized_grid, dtype=torch.bool))
 
-    def _create_grid(self):
-        grid = torch.zeros(self.grid_size, len(self.grid_bounds))
-        for i in range(len(self.grid_bounds)):
-            grid_diff = float(self.grid_bounds[i][1] - self.grid_bounds[i][0]) / (self.grid_size - 2)
-            grid[:, i] = torch.linspace(
-                self.grid_bounds[i][0] - grid_diff, self.grid_bounds[i][1] + grid_diff, self.grid_size
-            )
-
-        return grid
-
     @property
     def _tight_grid_bounds(self):
-        grid_spacings = tuple((bound[1] - bound[0]) / self.grid_size for bound in self.grid_bounds)
+        grid_spacings = tuple((bound[1] - bound[0]) / self.grid_sizes[i] for i, bound in enumerate(self.grid_bounds))
         return tuple(
             (bound[0] + 2.01 * spacing, bound[1] - 2.01 * spacing)
             for bound, spacing in zip(self.grid_bounds, grid_spacings)
@@ -150,12 +159,14 @@ class GridInterpolationKernel(GridKernel):
 
             # Update the grid if needed
             if update_grid:
-                grid_spacings = tuple((x_max - x_min) / (self.grid_size - 4.02) for x_min, x_max in zip(x_mins, x_maxs))
+                grid_spacings = tuple(
+                    (x_max - x_min) / (gs - 4.02) for gs, x_min, x_max in zip(self.grid_sizes, x_mins, x_maxs)
+                )
                 self.grid_bounds = tuple(
                     (x_min - 2.01 * spacing, x_max + 2.01 * spacing)
                     for x_min, x_max, spacing in zip(x_mins, x_maxs, grid_spacings)
                 )
-                grid = self._create_grid()
+                grid = create_grid(self.grid_sizes, self.grid_bounds)
                 self.update_grid(grid)
 
         base_lazy_tsr = self._inducing_forward(last_dim_is_batch=last_dim_is_batch, **params)
@@ -186,6 +197,3 @@ class GridInterpolationKernel(GridKernel):
 
     def prediction_strategy(self, train_inputs, train_prior_dist, train_labels, likelihood):
         return InterpolatedPredictionStrategy(train_inputs, train_prior_dist, train_labels, likelihood)
-
-    def num_outputs_per_input(self, x1, x2):
-        return self.base_kernel.num_outputs_per_input(x1, x2)
diff --git a/gpytorch/kernels/grid_kernel.py b/gpytorch/kernels/grid_kernel.py
index 3bf530b2..31342090 100644
--- a/gpytorch/kernels/grid_kernel.py
+++ b/gpytorch/kernels/grid_kernel.py
@@ -1,10 +1,13 @@
 #!/usr/bin/env python3
 
 import torch
+from torch import Tensor
 from .kernel import Kernel
 from ..lazy import delazify, ToeplitzLazyTensor, KroneckerProductLazyTensor
 from .. import settings
-from gpytorch.utils.grid import create_data_from_grid
+import gpytorch
+from gpytorch.utils.grid import create_data_from_grid, convert_legacy_grid
+from typing import List
 
 
 class GridKernel(Kernel):
@@ -23,42 +26,77 @@ class GridKernel(Kernel):
     Args:
         :attr:`base_kernel` (Kernel):
             The kernel to speed up with grid methods.
+        :attr:`grid` (Tensor):
+            A g x d tensor where column i consists of the projections of the
+            grid in dimension i.
         :attr:`active_dims` (tuple of ints, optional):
             Passed down to the `base_kernel`.
+        :attr:`interpolation_mode` (bool):
+            Used for GridInterpolationKernel where we want the covariance
+            between points in the projections of the grid of each dimension.
+            We do this by treating `grid` as d batches of g x 1 tensors by
+            calling base_kernel(grid, grid) with last_dim_is_batch to get a d x g x g Tensor
+            which we Kronecker product to get a g x g KroneckerProductLazyTensor.
 
     .. _Fast kernel learning for multidimensional pattern extrapolation:
         http://www.cs.cmu.edu/~andrewgw/manet.pdf
     """
 
-    def __init__(self, base_kernel, grid, interpolation_mode=False, active_dims=None):
+    def __init__(
+        self, base_kernel: Kernel, grid: List[Tensor], interpolation_mode: bool = False, active_dims: bool = None
+    ):
         super(GridKernel, self).__init__(active_dims=active_dims)
+        if torch.is_tensor(grid):
+            grid = convert_legacy_grid(grid)
         self.interpolation_mode = interpolation_mode
         self.base_kernel = base_kernel
-        self.register_buffer("grid", grid)
+        self.num_dims = len(grid)
+        self.register_buffer_list("grid", grid)
         if not self.interpolation_mode:
             self.register_buffer("full_grid", create_data_from_grid(grid))
 
+    def register_buffer_list(self, base_name, tensors):
+        """Helper to register several buffers at once under a single base name"""
+        for i, tensor in enumerate(tensors):
+            self.register_buffer(base_name + "_" + str(i), tensor)
+
     def train(self, mode=True):
         if hasattr(self, "_cached_kernel_mat"):
             del self._cached_kernel_mat
         return super(GridKernel, self).train(mode)
 
+    @property
+    def grid(self):
+        return [getattr(self, f"grid_{i}") for i in range(self.num_dims)]
+
     def update_grid(self, grid):
         """
         Supply a new `grid` if it ever changes.
         """
-        self.grid.detach().resize_(grid.size()).copy_(grid)
+        if torch.is_tensor(grid):
+            grid = convert_legacy_grid(grid)
+
+        if len(grid) != self.num_dims:
+            raise RuntimeError("New grid should have the same number of dimensions as before.")
+
+        for i in range(self.num_dims):
+            setattr(self, f"grid_{i}", grid[i])
 
         if not self.interpolation_mode:
-            full_grid = create_data_from_grid(self.grid)
-            self.full_grid.detach().resize_(full_grid).copy_(full_grid)
+            self.full_grid = create_data_from_grid(self.grid)
 
         if hasattr(self, "_cached_kernel_mat"):
             del self._cached_kernel_mat
         return self
 
+    @property
+    def is_ragged(self):
+        return not all(self.grid[0].size() == proj.size() for proj in self.grid)
+
     def forward(self, x1, x2, diag=False, last_dim_is_batch=False, **params):
         grid = self.grid
+        if last_dim_is_batch and self.is_ragged:
+            raise ValueError("last_dim_is_batch requires all dimensions to have same number of grid points")
 
         if not self.interpolation_mode:
             if len(x1.shape[:-2]):
@@ -69,28 +107,35 @@ class GridKernel(Kernel):
         if self.interpolation_mode or (torch.equal(x1, full_grid) and torch.equal(x2, full_grid)):
             if not self.training and hasattr(self, "_cached_kernel_mat"):
                 return self._cached_kernel_mat
-
-            n_dim = x1.size(-1)
-
+            # Can exploit Toeplitz structure if grid points in each dimension are equally
+            # spaced and using a translation-invariant kernel
             if settings.use_toeplitz.on():
-                first_item = grid[0:1]
-                covar_columns = self.base_kernel(first_item, grid, diag=False, last_dim_is_batch=True, **params)
-                covar_columns = delazify(covar_columns).squeeze(-2)
+                first_grid_point = [proj[0].unsqueeze(0) for proj in grid]
+                covars = [
+                    self.base_kernel(first, proj, last_dim_is_batch=False, **params)
+                    for first, proj in zip(first_grid_point, grid)
+                ]  # Each entry i contains a 1 x grid_size[i] covariance matrix
+                covars = [delazify(c) for c in covars]
+
                 if last_dim_is_batch:
-                    covars = [ToeplitzLazyTensor(covar_columns.squeeze(-2))]
+                    # Toeplitz expects batches of columns so we concatenate the
+                    # 1 x grid_size[i] tensors together
+                    # Note that this requires all the dimensions to have the same number of grid points
+                    covar = ToeplitzLazyTensor(torch.cat(covars, dim=-2))
                 else:
-                    covars = [ToeplitzLazyTensor(covar_columns[i : i + 1].squeeze(-2)) for i in range(n_dim)]
+                    # Non-batched ToeplitzLazyTensor expects a 1D tensor, so we squeeze out the row dimension
+                    covars = [ToeplitzLazyTensor(c.squeeze(-2)) for c in covars]
+                    # Due to legacy reasons, KroneckerProductLazyTensor(A, B, C) is actually (C Kron B Kron A)
+                    covar = KroneckerProductLazyTensor(*covars[::-1])
             else:
-                full_covar = self.base_kernel(grid, grid, last_dim_is_batch=True, **params)
+                covars = [
+                    self.base_kernel(proj, proj, last_dim_is_batch=False, **params) for proj in grid
+                ]  # Each entry i contains a grid_size[i] x grid_size[i] covariance matrix
                 if last_dim_is_batch:
-                    covars = [full_covar]
+                    # Note that this requires all the dimensions to have the same number of grid points
+                    covar = gpytorch.cat([c.unsqueeze(-3) for c in covars], dim=-3)
                 else:
-                    covars = [full_covar[i : i + 1].squeeze(0) for i in range(n_dim)]
-
-            if len(covars) > 1:
-                covar = KroneckerProductLazyTensor(*covars[::-1])
-            else:
-                covar = covars[0]
+                    covar = KroneckerProductLazyTensor(*covars[::-1])
 
             if not self.training:
                 self._cached_kernel_mat = covar
diff --git a/gpytorch/lazy/toeplitz_lazy_tensor.py b/gpytorch/lazy/toeplitz_lazy_tensor.py
index 1414b388..60137958 100644
--- a/gpytorch/lazy/toeplitz_lazy_tensor.py
+++ b/gpytorch/lazy/toeplitz_lazy_tensor.py
@@ -7,6 +7,14 @@ from ..utils.toeplitz import sym_toeplitz_matmul, sym_toeplitz_derivative_quadra
 
 class ToeplitzLazyTensor(LazyTensor):
     def __init__(self, column):
+        """
+        Args:
+            :attr: `column` (Tensor)
+                If `column` is a 1D Tensor of length `n`, this represents a
+                Toeplitz matrix with `column` as its first column.
+                If `column` is `b_1 x b_2 x ... x b_k x n`, then this represents a batch
+                `b_1 x b_2 x ... x b_k` of Toeplitz matrices.
+        """
         super(ToeplitzLazyTensor, self).__init__(column)
         self.column = column
 
@@ -35,7 +43,7 @@ class ToeplitzLazyTensor(LazyTensor):
         if res.dim() > self.column.dim():
             res = res.view(-1, *self.column.shape).sum(0)
 
-        return res,
+        return (res,)
 
     def _size(self):
         return torch.Size((*self.column.shape, self.column.size(-1)))
diff --git a/gpytorch/utils/grid.py b/gpytorch/utils/grid.py
index a340f67a..cb3d836e 100644
--- a/gpytorch/utils/grid.py
+++ b/gpytorch/utils/grid.py
@@ -2,6 +2,7 @@
 
 import math
 import torch
+from typing import List, Tuple
 
 
 def scale_to_bounds(x, lower_bound, upper_bound):
@@ -47,19 +48,55 @@ def choose_grid_size(train_inputs, ratio=1.0, kronecker_structure=True):
     if kronecker_structure:
         return int(ratio * math.pow(num_data, 1.0 / num_dim))
     else:
-        return (ratio * num_data)
+        return ratio * num_data
 
 
-def create_data_from_grid(grid):
-    grid_size = grid.size(-2)
-    grid_dim = grid.size(-1)
-    grid_data = torch.zeros(int(pow(grid_size, grid_dim)), grid_dim, device=grid.device)
-    prev_points = None
-    for i in range(grid_dim):
-        for j in range(grid_size):
-            grid_data[j * grid_size ** i : (j + 1) * grid_size ** i, i].fill_(grid[j, i])
-            if prev_points is not None:
-                grid_data[j * grid_size ** i : (j + 1) * grid_size ** i, :i].copy_(prev_points)
-        prev_points = grid_data[: grid_size ** (i + 1), : (i + 1)]
+def convert_legacy_grid(grid: torch.Tensor) -> List[torch.Tensor]:
+    return [grid[:, i] for i in range(grid.size(-1))]
 
+
+def create_data_from_grid(grid: List[torch.Tensor]) -> torch.Tensor:
+    """
+    Args:
+        :attr:`grid` (List[Tensor])
+            Each Tensor is a 1D set of increments for the grid in that dimension
+    Returns:
+        `grid_data` (Tensor)
+            Returns the set of points on the grid going by column-major order
+            (due to legacy reasons).
+    """
+    if torch.is_tensor(grid):
+        grid = convert_legacy_grid(grid)
+    ndims = len(grid)
+    assert all(axis.dim() == 1 for axis in grid)
+    projections = torch.meshgrid(*grid)
+    grid_tensor = torch.stack(projections, axis=-1)
+    # Note that if we did
+    #     grid_data = grid_tensor.reshape(-1, ndims)
+    # instead, we would be iterating through the points of our grid from the
+    # last data dimension to the first data dimension. However, due to legacy
+    # reasons, we need to iterate from the first data dimension to the last data
+    # dimension when creating grid_data
+    grid_data = grid_tensor.permute(*(reversed(range(ndims + 1)))).reshape(ndims, -1).transpose(0, 1)
     return grid_data
+
+
+def create_grid(
+    grid_sizes: List[int], grid_bounds: List[Tuple[float, float]], extend: bool = True
+) -> List[torch.Tensor]:
+    """
+    Creates a grid represented by a list of 1D Tensors representing the
+    projections of the grid into each dimension
+
+    If `extend`, we extend the grid by two points past the specified boundary
+    which can be important for getting good grid interpolations
+    """
+    grid = []
+    for i in range(len(grid_bounds)):
+        grid_diff = float(grid_bounds[i][1] - grid_bounds[i][0]) / (grid_sizes[i] - 2)
+        if extend:
+            proj = torch.linspace(grid_bounds[i][0] - grid_diff, grid_bounds[i][1] + grid_diff, grid_sizes[i])
+        else:
+            proj = torch.linspace(grid_bounds[i][0], grid_bounds[i][1], grid_sizes[i])
+        grid.append(proj)
+    return grid
diff --git a/gpytorch/utils/interpolation.py b/gpytorch/utils/interpolation.py
index 36e4277e..568e6087 100644
--- a/gpytorch/utils/interpolation.py
+++ b/gpytorch/utils/interpolation.py
@@ -2,12 +2,13 @@
 
 import torch
 from .broadcasting import _matmul_broadcast_shape
+from functools import reduce
+from operator import mul
+from .grid import convert_legacy_grid
+from typing import List
 
 
 class Interpolation(object):
-    """
-    """
-
     def _cubic_interpolation_kernel(self, scaled_grid_dist):
         """
         Computes the interpolation kernel u() for points X given the scaled
@@ -36,12 +37,20 @@ class Interpolation(object):
         res = res + (((-0.5 * U + 2.5).mul(U) - 4).mul(U) + 2) * U_ge_1_le_2
         return res
 
-    def interpolate(self, x_grid, x_target, interp_points=range(-2, 2)):
-        # Do some boundary checking
-        grid_mins = x_grid.min(0)[0]
-        grid_maxs = x_grid.max(0)[0]
+    def interpolate(self, x_grid: List[torch.Tensor], x_target: torch.Tensor, interp_points=range(-2, 2)):
+        if torch.is_tensor(x_grid):
+            x_grid = convert_legacy_grid(x_grid)
+        num_target_points = x_target.size(0)
+        num_dim = x_target.size(-1)
+        assert num_dim == len(x_grid)
+
+        grid_sizes = [len(x_grid[i]) for i in range(num_dim)]
+        # Do some boundary checking, # min/max along each dimension
+        x_target_max = x_target.max(0)[0]
         x_target_min = x_target.min(0)[0]
-        x_target_max = x_target.min(0)[0]
+        grid_mins = torch.stack([x_grid[i].min() for i in range(num_dim)], dim=0).to(x_target_min)
+        grid_maxs = torch.stack([x_grid[i].max() for i in range(num_dim)], dim=0).to(x_target_max)
+
         lt_min_mask = (x_target_min - grid_mins).lt(-1e-7)
         gt_max_mask = (x_target_max - grid_maxs).gt(1e-7)
         if lt_min_mask.sum().item():
@@ -74,31 +83,32 @@ class Interpolation(object):
             )
 
         # Now do interpolation
-        interp_points = torch.tensor(interp_points, dtype=x_grid.dtype, device=x_grid.device)
-        interp_points_flip = interp_points.flip(0)
+        interp_points = torch.tensor(interp_points, dtype=x_grid[0].dtype, device=x_grid[0].device)
+        interp_points_flip = interp_points.flip(0)  # [1, 0, -1, -2]
 
-        num_grid_points = x_grid.size(0)
-        num_target_points = x_target.size(0)
-        num_dim = x_target.size(-1)
         num_coefficients = len(interp_points)
 
         interp_values = torch.ones(
-            num_target_points, num_coefficients ** num_dim, dtype=x_grid.dtype, device=x_grid.device
+            num_target_points, num_coefficients ** num_dim, dtype=x_grid[0].dtype, device=x_grid[0].device
         )
         interp_indices = torch.zeros(
-            num_target_points, num_coefficients ** num_dim, dtype=torch.long, device=x_grid.device
+            num_target_points, num_coefficients ** num_dim, dtype=torch.long, device=x_grid[0].device
         )
 
         for i in range(num_dim):
-            grid_delta = x_grid[1, i] - x_grid[0, i]
-            lower_grid_pt_idxs = torch.floor((x_target[:, i] - x_grid[0, i]) / grid_delta).squeeze()
-            lower_pt_rel_dists = (x_target[:, i] - x_grid[0, i]) / grid_delta - lower_grid_pt_idxs
-            lower_grid_pt_idxs = lower_grid_pt_idxs - interp_points.max()
+            num_grid_points = x_grid[i].size(0)
+            grid_delta = x_grid[i][1] - x_grid[i][0]
+            # left-bounding grid point in index space
+            lower_grid_pt_idxs = torch.floor((x_target[:, i] - x_grid[i][0]) / grid_delta)
+            # distance from that left-bounding grid point, again in index space
+            lower_pt_rel_dists = (x_target[:, i] - x_grid[i][0]) / grid_delta - lower_grid_pt_idxs
+            lower_grid_pt_idxs = lower_grid_pt_idxs - interp_points.max()  # ends up being the left-most (relevant) pt
             lower_grid_pt_idxs.detach_()
 
             if len(lower_grid_pt_idxs.shape) == 0:
                 lower_grid_pt_idxs = lower_grid_pt_idxs.unsqueeze(0)
 
+            # get the interp. coeff. based on distances to interpolating points
             scaled_dist = lower_pt_rel_dists.unsqueeze(-1) + interp_points_flip.unsqueeze(-2)
             dim_interp_values = self._cubic_interpolation_kernel(scaled_dist)
 
@@ -109,7 +119,7 @@ class Interpolation(object):
 
             if num_left > 0:
                 left_boundary_pts.squeeze_(1)
-                x_grid_first = x_grid[:num_coefficients, i].unsqueeze(1).t().expand(num_left, num_coefficients)
+                x_grid_first = x_grid[i][:num_coefficients].unsqueeze(1).t().expand(num_left, num_coefficients)
 
                 grid_targets = x_target.select(1, i)[left_boundary_pts].unsqueeze(1).expand(num_left, num_coefficients)
                 dists = torch.abs(x_grid_first - grid_targets)
@@ -125,7 +135,7 @@ class Interpolation(object):
 
             if num_right > 0:
                 right_boundary_pts.squeeze_(1)
-                x_grid_last = x_grid[-num_coefficients:, i].unsqueeze(1).t().expand(num_right, num_coefficients)
+                x_grid_last = x_grid[i][-num_coefficients:].unsqueeze(1).t().expand(num_right, num_coefficients)
 
                 grid_targets = x_target.select(1, i)[right_boundary_pts].unsqueeze(1)
                 grid_targets = grid_targets.expand(num_right, num_coefficients)
@@ -138,13 +148,15 @@ class Interpolation(object):
                     lower_grid_pt_idxs[right_boundary_pts[j]] = num_grid_points - num_coefficients
 
             offset = (interp_points - interp_points.min()).long().unsqueeze(-2)
-            dim_interp_indices = lower_grid_pt_idxs.long().unsqueeze(-1) + offset
+            dim_interp_indices = lower_grid_pt_idxs.long().unsqueeze(-1) + offset  # indices of corresponding ind. pts.
 
             n_inner_repeat = num_coefficients ** i
             n_outer_repeat = num_coefficients ** (num_dim - i - 1)
-            index_coeff = num_grid_points ** (num_dim - i - 1)
+            # index_coeff = num_grid_points ** (num_dim - i - 1)  # TODO: double check
+            index_coeff = reduce(mul, grid_sizes[i + 1 :], 1)  # Think this is right...
             dim_interp_indices = dim_interp_indices.unsqueeze(-1).repeat(1, n_inner_repeat, n_outer_repeat)
             dim_interp_values = dim_interp_values.unsqueeze(-1).repeat(1, n_inner_repeat, n_outer_repeat)
+            # compute the lexicographical position of the indices in the d-dimensional grid points
             interp_indices = interp_indices.add(dim_interp_indices.view(num_target_points, -1).mul(index_coeff))
             interp_values = interp_values.mul(dim_interp_values.view(num_target_points, -1))
 
@@ -187,7 +199,7 @@ def left_t_interp(interp_indices, interp_values, rhs, output_dim):
 
     # Multiply the rhs by the interp_values
     # This multiplication here will give us the ability to perform backprop
-    values = (rhs.unsqueeze(-2) * interp_values.unsqueeze(-1))
+    values = rhs.unsqueeze(-2) * interp_values.unsqueeze(-1)
 
     # Define a bunch of sizes
     num_data, num_interp = interp_values.shape[-2:]


Test patch:
diff --git a/test/kernels/test_grid_kernel.py b/test/kernels/test_grid_kernel.py
index f519a976..df8c7057 100644
--- a/test/kernels/test_grid_kernel.py
+++ b/test/kernels/test_grid_kernel.py
@@ -2,22 +2,14 @@
 
 import torch
 import unittest
-from gpytorch.kernels import RBFKernel, GridKernel, GridInterpolationKernel
+from gpytorch.kernels import RBFKernel, GridKernel
 from gpytorch.lazy import KroneckerProductLazyTensor
+from gpytorch.utils.grid import create_data_from_grid
 
-cv = GridInterpolationKernel(RBFKernel(), grid_size=10, grid_bounds=[(0, 1), (0, 2)])
-grid = cv.grid
 
-grid_size = grid.size(-2)
-grid_dim = grid.size(-1)
-grid_data = torch.zeros(int(pow(grid_size, grid_dim)), grid_dim)
-prev_points = None
-for i in range(grid_dim):
-    for j in range(grid_size):
-        grid_data[j * grid_size ** i : (j + 1) * grid_size ** i, i].fill_(grid[j, i])
-        if prev_points is not None:
-            grid_data[j * grid_size ** i : (j + 1) * grid_size ** i, :i].copy_(prev_points)
-    prev_points = grid_data[: grid_size ** (i + 1), : (i + 1)]
+grid = [torch.linspace(0, 1, 5), torch.linspace(0, 2, 3)]
+d = len(grid)
+grid_data = create_data_from_grid(grid)
 
 
 class TestGridKernel(unittest.TestCase):
@@ -32,7 +24,7 @@ class TestGridKernel(unittest.TestCase):
 
     def test_nongrid_grid(self):
         base_kernel = RBFKernel()
-        data = torch.randn(5, 2)
+        data = torch.randn(5, d)
         kernel = GridKernel(base_kernel, grid)
         grid_eval = kernel(grid_data, data).evaluate()
         actual_eval = base_kernel(grid_data, data).evaluate()
@@ -40,7 +32,7 @@ class TestGridKernel(unittest.TestCase):
 
     def test_nongrid_nongrid(self):
         base_kernel = RBFKernel()
-        data = torch.randn(5, 2)
+        data = torch.randn(5, d)
         kernel = GridKernel(base_kernel, grid)
         grid_eval = kernel(data, data).evaluate()
         actual_eval = base_kernel(data, data).evaluate()