# Buggy code
def __init__(
    self,
    base_lazy_tensor,
    left_interp_indices=None,
    left_interp_values=None,
    right_interp_indices=None,
    right_interp_values=None,
):
    base_lazy_tensor = lazify(base_lazy_tensor)

    if left_interp_indices is None:
        num_rows = base_lazy_tensor.size(-2)
        left_interp_indices = torch.arange(
            0, num_rows, dtype=torch.long, device=base_lazy_tensor.device
        )
        left_interp_indices.unsqueeze_(-1)
        left_interp_indices = left_interp_indices.expand(
            *base_lazy_tensor.batch_shape, num_rows, 1
        )

    if left_interp_values is None:
        left_interp_values = torch.ones(
            left_interp_indices.size(),
            dtype=base_lazy_tensor.dtype,
            device=base_lazy_tensor.device,
        )
    else:
        if left_interp_indices.size() != left_interp_values.size():
            raise RuntimeError(
                "Expected left_interp_indices ({}) to have the same size as "
                "left_interp_values ({})".format(
                    left_interp_indices.size(), left_interp_values.size()
                )
            )

    if right_interp_indices is None:
        num_rows = base_lazy_tensor.size(-2)
        right_interp_indices = torch.arange(
            0, num_rows, dtype=torch.long, device=base_lazy_tensor.device
        )
        right_interp_indices.unsqueeze_(-1)
        right_interp_indices = right_interp_indices.expand(
            *base_lazy_tensor.batch_shape, num_rows, 1
        )

    if right_interp_values is None:
        right_interp_values = torch.ones(
            right_interp_indices.size(),
            dtype=base_lazy_tensor.dtype,
            device=base_lazy_tensor.device,
        )
    else:
        if left_interp_indices.size() != left_interp_values.size():
            raise RuntimeError(
                "Expected left_interp_indices ({}) to have the same size as "
                "left_interp_values ({})".format(
                    left_interp_indices.size(), left_interp_values.size()
                )
            )

    # Make sure that left/right interp tensors have the same batch shape as the base_lazy_tensor
    if left_interp_indices.shape[:-2] != base_lazy_tensor.batch_shape:
        raise RuntimeError(
            "left interp size ({}) is incompatible with base_lazy_tensor size ({}). Make sure the two "
            "have the same number of batch dimensions".format(
                left_interp_indices.size(), base_lazy_tensor.size()
            )
        )
    if right_interp_indices.shape[:-2] != base_lazy_tensor.batch_shape:
        raise RuntimeError(
            "right interp size ({}) is incompatible with base_lazy_tensor size ({}). Make sure the two "
            "have the same number of batch dimensions".format(
                right_interp_indices.size(), base_lazy_tensor.size()
            )
        )

    super(InterpolatedLazyTensor, self).__init__(
        base_lazy_tensor,
        left_interp_indices,
        left_interp_values,
        right_interp_indices,
        right_interp_values,
    )
    self.base_lazy_tensor = base_lazy_tensor
    self.left_interp_indices = left_interp_indices
    self.left_interp_values = left_interp_values
    self.right_interp_indices = right_interp_indices
    self.right_interp_values = right_interp_values


# Traceback (RuntimeError)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-62-28172f8b7beb> in <module>()
1 with torch.no_grad(), gpytorch.settings.fast_pred_var():
----> 2     observed_pred_y1 = likelihood(model(test_x, tast_i_task1))

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
199             )
200
--> 201             full_output = super(ExactGP, self).__call__(*full_inputs, **kwargs)
202             if settings.debug().on():
203                 if not isinstance(full_output, MultivariateNormal):

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/module.py in __call__(self, *inputs, **kwargs)
18
19     def __call__(self, *inputs, **kwargs):
---> 20         outputs = self.forward(*inputs, **kwargs)
21         if isinstance(outputs, list):
22             return [_validate_module_outputs(output) for output in outputs]

<ipython-input-52-b565681c3db1> in forward(self, x, i)
17         covar_i = self.task_covar_module(i)
18         # Multiply the two together to get the covariance we want
---> 19         covar = covar_x.mul(covar_i)
20
21         return gpytorch.distributions.MultivariateNormal(mean_x, covar)

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/lazy/lazy_tensor.py in mul(self, other)
950                 return self._mul_constant(other.view(*other.shape[:-2]))
951
--> 952         return self._mul_matrix(other)
953
954     def ndimension(self):

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/lazy/lazy_tensor.py in _mul_matrix(self, other)
443         """
444         from .mul_lazy_tensor import MulLazyTensor
--> 445         return MulLazyTensor(self, other).evaluate_kernel()
446
447     def _preconditioner(self):

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/lazy/lazy_tensor.py in evaluate_kernel(self)
734         all lazily evaluated kernels actually evaluated.
735         """
--> 736         return self.representation_tree()(*self.representation())
737
738     def inv_matmul(self, right_tensor, left_tensor=None):

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/lazy/mul_lazy_tensor.py in representation_tree(self)
235
236     def representation_tree(self):
--> 237         if self.non_lazy_self is not None:
238             return self.non_lazy_self.representation_tree()
239         else:

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/lazy/mul_lazy_tensor.py in non_lazy_self(self)
37         if hasattr(self, "_non_lazy_self"):
38             return self._non_lazy_self[0]
---> 39         elif len(self._args) == 1:
40             return self._args[0]
41         else:

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/lazy/mul_lazy_tensor.py in _args(self)
54         if not hasattr(self, "_mul_args_memo") and not hasattr(self, "_non_lazy_self"):
55             lazy_tensors = sorted(
---> 56                 (lv.evaluate_kernel() for lv in self.lazy_tensors), key=lambda lv: lv.root_decomposition_size()
57             )
58

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/lazy/mul_lazy_tensor.py in <genexpr>(.0)
54         if not hasattr(self, "_mul_args_memo") and not hasattr(self, "_non_lazy_self"):
55             lazy_tensors = sorted(
---> 56                 (lv.evaluate_kernel() for lv in self.lazy_tensors), key=lambda lv: lv.root_decomposition_size()
57             )
58

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate_kernel(self)
159             with settings.lazily_evaluate_kernels(False):
160                 self._cached_kernel_eval = self.kernel(
--> 161                     x1, x2, diag=False, batch_dims=self.batch_dims, **self.params
162                 )
163

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/kernels/kernel.py in __call__(self, x1, x2, diag, batch_dims, **params)
396                 res = LazyEvaluatedKernelTensor(self, x1_, x2_, batch_dims=batch_dims, **params)
397             else:
--> 398                 res = super(Kernel, self).__call__(x1_, x2_, batch_dims=batch_dims, **params)
399
400             # TODO: remove bach checking once kernels support arbitrary batch dimensions

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/module.py in __call__(self, *inputs, **kwargs)
18
19     def __call__(self, *inputs, **kwargs):
---> 20         outputs = self.forward(*inputs, **kwargs)
21         if isinstance(outputs, list):
22             return [_validate_module_outputs(output) for output in outputs]

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/kernels/index_kernel.py in forward(self, i1, i2, **params)
78     def forward(self, i1, i2, **params):
79         covar_matrix = self._eval_covar_matrix()
---> 80         res = InterpolatedLazyTensor(base_lazy_tensor=covar_matrix, left_interp_indices=i1, right_interp_indices=i2)
81         return res

/data/users/balandat/fbsource/fbcode/buck-out/dev/gen/bento/kernels/bento_kernel_ae_lazarus#link-tree/gpytorch/lazy/interpolated_lazy_tensor.py in __init__(self, base_lazy_tensor, left_interp_indices, left_interp_values, right_interp_indices, right_interp_values)
60             raise RuntimeError(
61                 "left interp size ({}) is incompatible with base_lazy_tensor size ({}). Make sure the two "
---> 62                 "have the same number of batch dimensions".format(left_interp_indices.size(), base_lazy_tensor.size())
63             )
64         if right_interp_indices.shape[:-2] != base_lazy_tensor.batch_shape:

RuntimeError: left interp size (torch.Size([2, 251, 1])) is incompatible with base_lazy_tensor size (torch.Size([1, 2, 2])). Make sure the two have the same number of batch dimensions

# Fixed code
def __init__(
    self,
    base_lazy_tensor,
    left_interp_indices=None,
    left_interp_values=None,
    right_interp_indices=None,
    right_interp_values=None,
):
    base_lazy_tensor = lazify(base_lazy_tensor)

    if left_interp_indices is None:
        num_rows = base_lazy_tensor.size(-2)
        left_interp_indices = torch.arange(
            0, num_rows, dtype=torch.long, device=base_lazy_tensor.device
        )
        left_interp_indices.unsqueeze_(-1)
        left_interp_indices = left_interp_indices.expand(
            *base_lazy_tensor.batch_shape, num_rows, 1
        )

    if left_interp_values is None:
        left_interp_values = torch.ones(
            left_interp_indices.size(),
            dtype=base_lazy_tensor.dtype,
            device=base_lazy_tensor.device,
        )

    if right_interp_indices is None:
        num_rows = base_lazy_tensor.size(-2)
        right_interp_indices = torch.arange(
            0, num_rows, dtype=torch.long, device=base_lazy_tensor.device
        )
        right_interp_indices.unsqueeze_(-1)
        right_interp_indices = right_interp_indices.expand(
            *base_lazy_tensor.batch_shape, num_rows, 1
        )

    if right_interp_values is None:
        right_interp_values = torch.ones(
            right_interp_indices.size(),
            dtype=base_lazy_tensor.dtype,
            device=base_lazy_tensor.device,
        )

    if left_interp_indices.shape[:-2] != base_lazy_tensor.batch_shape:
        try:
            base_lazy_tensor = base_lazy_tensor._expand_batch(
                left_interp_indices.shape[:-2]
            )
        except RuntimeError:
            raise RuntimeError(
                "interp size ({}) is incompatible with base_lazy_tensor size ({}). ".format(
                    right_interp_indices.size(), base_lazy_tensor.size()
                )
            )

    super(InterpolatedLazyTensor, self).__init__(
        base_lazy_tensor,
        left_interp_indices,
        left_interp_values,
        right_interp_indices,
        right_interp_values,
    )
    self.base_lazy_tensor = base_lazy_tensor
    self.left_interp_indices = left_interp_indices
    self.left_interp_values = left_interp_values
    self.right_interp_indices = right_interp_indices
    self.right_interp_values = right_interp_values