Add support for batches in the index kernel
@@ -39,14 +39,7 @@ class IndexKernel(Kernel):
The element-wise log of the :math:`\mathbf v` vector.
"""
- def __init__(
- self,
- num_tasks,
- rank=1,
- prior=None,
- var_constraint=None,
- **kwargs
- ):
+ def __init__(self, num_tasks, rank=1, prior=None, var_constraint=None, **kwargs):
if rank > num_tasks:
raise RuntimeError("Cannot create a task covariance matrix larger than the number of tasks")
super().__init__(**kwargs)
@@ -76,7 +69,12 @@ class IndexKernel(Kernel):
def _eval_covar_matrix(self):
var = self.var
- D = var * torch.eye(var.shape[-1], dtype=var.dtype, device=var.device)
+ eye = torch.eye(var.shape[-1], dtype=var.dtype, device=var.device)
+ if len(self.batch_shape) > 0:
+ eye = eye.view((1, var.shape[-1], var.shape[-1])).repeat(*self.batch_shape, 1, 1)
+ D = var.repeat(1, var.shape[-1]).view(*self.batch_shape, var.shape[-1], var.shape[-1]) * eye
+ else:
+ D = var * eye
return self.covar_factor.matmul(self.covar_factor.transpose(-1, -2)) + D
@property
@@ -87,5 +85,9 @@ class IndexKernel(Kernel):
def forward(self, i1, i2, **params):
covar_matrix = self._eval_covar_matrix()
- res = InterpolatedLazyTensor(base_lazy_tensor=covar_matrix, left_interp_indices=i1, right_interp_indices=i2)
+ res = InterpolatedLazyTensor(
+ base_lazy_tensor=covar_matrix,
+ left_interp_indices=i1.repeat(*self.batch_shape, 1, 1),
+ right_interp_indices=i2.repeat(*self.batch_shape, 1, 1),
+ )
return res
