Refactor ConstantMulLazyTensor.__getitem__
import torch
from .lazy_tensor import LazyTensor
+from .. import settings
class ConstantMulLazyTensor(LazyTensor):
@@ -128,20 +129,42 @@ class ConstantMulLazyTensor(LazyTensor):
constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
return res * constant
- def __getitem__(self, i):
- # TODO: Fix to support ellipsis
- constant = self.constant
- if constant.numel() > 1:
- first_index = i[0] if isinstance(i, tuple) else i
- constant = constant[first_index]
+ def __getitem__(self, index):
+ # NOTE TO FUTURE SELF:
+ # This custom __getitem__ is actually very important!
+ # It prevents constructing an InterpolatedLazyTensor when one isn't needed
+ # This effects runntimes by up to 5x on simple exat GPs
+
+ # Process the index
+ ndimension = self.ndimension()
+ index = list(index) if isinstance(index, tuple) else [index]
+ index = [torch.tensor(idx) if isinstance(idx, list) else idx for idx in index]
+
+ # Handle the ellipsis
+ # Find the index of the ellipsis
+ ellipsis_locs = [index for index, item in enumerate(index) if item is Ellipsis]
+ if settings.debug.on():
+ if len(ellipsis_locs) > 1:
+ raise RuntimeError(
+ "Cannot have multiple ellipsis in a __getitem__ call. LazyTensor {} "
+ " received index {}.".format(self, index)
+ )
+ if len(ellipsis_locs) == 1:
+ ellipsis_loc = ellipsis_locs[0]
+ num_to_fill_in = ndimension - (len(index) - 1)
+ index = index[:ellipsis_loc] + [slice(None, None, None)] * num_to_fill_in + index[ellipsis_loc + 1:]
- base_lazy_tensor = self.base_lazy_tensor.__getitem__(i)
+ # Pad the index with empty slices
+ index += [slice(None, None, None)] * (ndimension - len(index))
- if torch.is_tensor(base_lazy_tensor):
- constant = constant.view(*constant.shape, *[1] * (len(base_lazy_tensor.shape) - len(constant.shape)))
- return constant * base_lazy_tensor
+ # Convert the index back to a tuple
+ index = tuple(index)
- if constant.numel() == 1:
- constant = constant.squeeze()
+ # Run __getitem__ on the base_lazy_tensor and the constant
+ base_lazy_tensor = self.base_lazy_tensor.__getitem__(index)
+ constant = self.constant[index[:-2]]
+
+ if torch.is_tensor(base_lazy_tensor):
+ constant = constant.view(*constant.shape, *[1] * (base_lazy_tensor.dim() - constant.dim()))
return base_lazy_tensor * constant
