Repository: cornellius-gp/gpytorch
Pull request #331 (2018-10-20 21:14:00)

Repeat lazy tensor
This adds a default repeat to LazyTensors


Patch:
diff --git a/gpytorch/functions/__init__.py b/gpytorch/functions/__init__.py
index 28862784..d4a4a75e 100644
--- a/gpytorch/functions/__init__.py
+++ b/gpytorch/functions/__init__.py
@@ -30,7 +30,7 @@ def add_diag(input, diag):
         return input.add_diag(diag)
 
 
-def add_jitter(mat):
+def add_jitter(mat, jitter_val=1e-3):
     """
     Adds "jitter" to the diagonal of a matrix.
     This ensures that a matrix that *should* be positive definite *is* positive definate.
@@ -40,9 +40,9 @@ def add_jitter(mat):
     Returns: (matrix nxn)
     """
     if hasattr(mat, "add_jitter"):
-        return mat.add_jitter()
+        return mat.add_jitter(jitter_val)
     else:
-        diag = torch.eye(mat.size(-1), dtype=mat.dtype, device=mat.device).mul_(1e-3)
+        diag = torch.eye(mat.size(-1), dtype=mat.dtype, device=mat.device).mul_(jitter_val)
         if mat.ndimension() == 3:
             return mat + diag.unsqueeze(0).expand(mat.size(0), mat.size(1), mat.size(2))
         else:
@@ -173,7 +173,7 @@ def inv_quad(mat, tensor):
     return res
 
 
-def inv_quad_log_det(mat, inv_quad_rhs=None, log_det=False):
+def inv_quad_log_det(mat, inv_quad_rhs=None, log_det=False, reduce_inv_quad=True):
     """
     Computes an inverse quadratic form (w.r.t mat) with several right hand sides.
     I.e. computes tr( tensor^T mat^{-1} tensor )
@@ -187,11 +187,11 @@ def inv_quad_log_det(mat, inv_quad_rhs=None, log_det=False):
         - scalar - log determinant
     """
     if hasattr(mat, "inv_quad_log_det"):
-        return mat.inv_quad_log_det(inv_quad_rhs, log_det)
+        return mat.inv_quad_log_det(inv_quad_rhs, log_det, reduce_inv_quad=reduce_inv_quad)
     else:
         from ..lazy.non_lazy_tensor import NonLazyTensor
 
-        return NonLazyTensor(mat).inv_quad_log_det(inv_quad_rhs, log_det)
+        return NonLazyTensor(mat).inv_quad_log_det(inv_quad_rhs, log_det, reduce_inv_quad=reduce_inv_quad)
 
 
 def log_det(mat):
diff --git a/gpytorch/functions/_inv_quad_log_det.py b/gpytorch/functions/_inv_quad_log_det.py
index 3d93c528..fc5e72c0 100644
--- a/gpytorch/functions/_inv_quad_log_det.py
+++ b/gpytorch/functions/_inv_quad_log_det.py
@@ -136,7 +136,7 @@ class InvQuadLogDet(Function):
         # Extract inv_quad solves from all the solves
         if self.inv_quad:
             inv_quad_solves = solves.narrow(-1, num_random_probes, num_inv_quad_solves)
-            inv_quad_term = (inv_quad_solves * inv_quad_rhs).sum(-1).sum(-1, keepdim=(self.batch_shape is None))
+            inv_quad_term = (inv_quad_solves * inv_quad_rhs).sum(-2)
 
         self.num_random_probes = num_random_probes
         self.num_inv_quad_solves = num_inv_quad_solves
@@ -170,13 +170,10 @@ class InvQuadLogDet(Function):
 
         # Fix grad_output sizes
         if self.inv_quad:
-            inv_quad_grad_output = inv_quad_grad_output.unsqueeze(-1)
-            if self.batch_shape is not None:
-                inv_quad_grad_output.unsqueeze_(-1)
+            inv_quad_grad_output = inv_quad_grad_output.unsqueeze(-2)
         if compute_log_det_grad:
             log_det_grad_output = log_det_grad_output.unsqueeze(-1)
-            if self.batch_shape is not None:
-                log_det_grad_output.unsqueeze_(-1)
+            log_det_grad_output.unsqueeze_(-1)
 
         # Divide up the solves
         probe_vector_solves = None
diff --git a/gpytorch/lazy/__init__.py b/gpytorch/lazy/__init__.py
index b8dd4729..8b035307 100644
--- a/gpytorch/lazy/__init__.py
+++ b/gpytorch/lazy/__init__.py
@@ -5,6 +5,7 @@ from __future__ import unicode_literals
 
 from .lazy_tensor import LazyTensor
 from .added_diag_lazy_tensor import AddedDiagLazyTensor
+from .batch_repeat_lazy_tensor import BatchRepeatLazyTensor
 from .block_lazy_tensor import BlockLazyTensor
 from .block_diag_lazy_tensor import BlockDiagLazyTensor
 from .chol_lazy_tensor import CholLazyTensor
@@ -28,6 +29,7 @@ __all__ = [
     "LazyTensor",
     "LazyEvaluatedKernelTensor",
     "AddedDiagLazyTensor",
+    "BatchRepeatLazyTensor",
     "BlockLazyTensor",
     "BlockDiagLazyTensor",
     "CholLazyTensor",
diff --git a/gpytorch/lazy/batch_repeat_lazy_tensor.py b/gpytorch/lazy/batch_repeat_lazy_tensor.py
new file mode 100644
index 00000000..e8ede2aa
--- /dev/null
+++ b/gpytorch/lazy/batch_repeat_lazy_tensor.py
@@ -0,0 +1,182 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import itertools
+import torch
+
+from .lazy_tensor import LazyTensor
+
+
+class BatchRepeatLazyTensor(LazyTensor):
+    def __init__(self, base_lazy_tensor, batch_repeat=torch.Size((1,))):
+        if not isinstance(batch_repeat, torch.Size):
+            raise RuntimeError(
+                "batch_repeat must be a torch.Size, got a {} instead".format(batch_repeat.__class__.__name__)
+            )
+
+        super(BatchRepeatLazyTensor, self).__init__(base_lazy_tensor, batch_repeat=batch_repeat)
+        self.base_lazy_tensor = base_lazy_tensor
+        self.batch_repeat = batch_repeat
+
+    def _batch_get_indices(self, batch_indices, row_indices, col_indices):
+        return self._get_indices(batch_indices, row_indices=row_indices, col_indices=col_indices)
+
+    def _get_indices(self, *batch_indices, row_indices, col_indices):
+        num_true_batch_dims = len(self.base_lazy_tensor.batch_shape)
+        batch_indices = [index % size for index, size in zip(batch_indices, self._padded_base_batch_shape)]
+        batch_indices = batch_indices[-num_true_batch_dims:] if num_true_batch_dims else []
+        if len(batch_indices) == 1:
+            return self.base_lazy_tensor._batch_get_indices(*batch_indices, row_indices, col_indices)
+        else:
+            return self.base_lazy_tensor._get_indices(*batch_indices, row_indices, col_indices)
+
+    def _matmul(self, rhs):
+        rhs = self._move_repeat_batches_to_columns(rhs)
+        res = self.base_lazy_tensor._matmul(rhs)
+        res = self._move_repeat_batches_back(res)
+        return res
+
+    def _move_repeat_batches_back(self, batch_matrix):
+        """
+        The opposite of _move_repeat_batches_to_columns
+
+        Takes a b x m x nr tensor, and moves the batches associated with repeating
+        So that the tensor is now rb x m x n.
+        """
+        orig_shape = (*self.batch_shape, batch_matrix.size(-2), -1)
+        padded_base_batch_shape = self._padded_base_batch_shape
+
+        # Now we have to move the columns back to their original repeat dimensions
+        batch_matrix = batch_matrix.view(*padded_base_batch_shape, batch_matrix.size(-2), -1, *self.batch_repeat)
+        dims = tuple(itertools.chain.from_iterable(
+            [i + len(orig_shape), i] for i in range(len(padded_base_batch_shape))
+        )) + (self.dim() - 2, self.dim() - 1)
+        batch_matrix = batch_matrix.permute(*dims).contiguous()
+
+        # Combine the repeat and the batch dimensions, and return the batch_matrixult!
+        batch_matrix = batch_matrix.view(*orig_shape)
+        return batch_matrix
+
+    def _move_repeat_batches_to_columns(self, batch_matrix):
+        """
+        Takes a rb x m x n tensor, and moves the batches associated with repeating
+        So that the tensor is now b x m x nr.
+        This allows us to use the base_lazy_tensor routines.
+        """
+        batch_matrix_shape = batch_matrix.shape
+        padded_base_batch_shape = self._padded_base_batch_shape
+
+        # Reshape batch_matrix so that each batch dimension is split in two:
+        # The repeated part, and the actual part
+        split_shape = torch.Size(tuple(itertools.chain.from_iterable(
+            [repeat, size] for repeat, size in zip(self.batch_repeat, padded_base_batch_shape)
+        )) + batch_matrix_shape[-2:])
+        batch_matrix = batch_matrix.view(*split_shape)
+
+        # Now chuck the repeat parts of the batch dimensions into the last dimension of batch_matrix
+        # These will act like extra columns of the batch matrix that we are multiplying against
+        # The repeated part, and the actual part
+        repeat_dims = range(0, len(self.batch_repeat) * 2, 2)
+        batch_dims = range(1, len(self.batch_repeat) * 2, 2)
+        batch_matrix = batch_matrix.permute(*batch_dims, -2, -1, *repeat_dims).contiguous()
+        batch_matrix = batch_matrix.view(*self.base_lazy_tensor.batch_shape, batch_matrix_shape[-2], -1)
+        return batch_matrix
+
+    @property
+    def _padded_base_batch_shape(self):
+        base_batch_shape = self.base_lazy_tensor.batch_shape
+        return torch.Size(
+            ([1] * (len(self.batch_repeat) - len(base_batch_shape))) + list(base_batch_shape)
+        )
+
+    def _quad_form_derivative(self, left_vectors, right_vectors):
+        left_vectors = self._move_repeat_batches_to_columns(left_vectors)
+        right_vectors = self._move_repeat_batches_to_columns(right_vectors)
+        return self.base_lazy_tensor._quad_form_derivative(left_vectors, right_vectors)
+
+    def _size(self):
+        repeated_batch_shape = torch.Size(
+            size * repeat for size, repeat in zip(self._padded_base_batch_shape, self.batch_repeat)
+        )
+        res = torch.Size(repeated_batch_shape + self.base_lazy_tensor.matrix_shape)
+        return res
+
+    def _transpose_nonbatch(self):
+        return self.__class__(self.base_lazy_tensor._transpose_nonbatch(), batch_repeat=self.batch_repeat)
+
+    def add_jitter(self, jitter_val=1e-3):
+        return self.__class__(self.base_lazy_tensor.add_jitter(jitter_val=jitter_val), batch_repeat=self.batch_repeat)
+
+    def inv_quad_log_det(self, inv_quad_rhs=None, log_det=False, reduce_inv_quad=True):
+        if not self.is_square:
+            raise RuntimeError(
+                "inv_quad_log_det only operates on (batches of) square (positive semi-definite) LazyTensors. "
+                "Got a {} of size {}.".format(self.__class__.__name__, self.size())
+            )
+
+        if inv_quad_rhs is not None:
+            if self.dim() != inv_quad_rhs.dim():
+                raise RuntimeError(
+                    "LazyTensor (size={}) and right-hand-side Tensor (size={}) should have the same number "
+                    "of dimensions.".format(self.shape, inv_quad_rhs.shape)
+                )
+            elif self.batch_shape != inv_quad_rhs.shape[:-2] or self.shape[-1] != inv_quad_rhs.shape[-2]:
+                raise RuntimeError(
+                    "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                        self.shape, inv_quad_rhs.shape
+                    )
+                )
+
+        if inv_quad_rhs is not None:
+            inv_quad_rhs = self._move_repeat_batches_to_columns(inv_quad_rhs)
+
+        inv_quad_term, log_det_term = self.base_lazy_tensor.inv_quad_log_det(
+            inv_quad_rhs, log_det, reduce_inv_quad=False
+        )
+
+        if inv_quad_term.numel():
+            inv_quad_term = inv_quad_term.view(*inv_quad_term.shape[:-1], -1, self.batch_repeat.numel())
+            inv_quad_term = self._move_repeat_batches_back(inv_quad_term).squeeze(-1)
+            if reduce_inv_quad:
+                inv_quad_term = inv_quad_term.sum(-1)
+
+        if log_det_term.numel():
+            log_det_term = log_det_term.repeat(*self.batch_repeat)
+
+        return inv_quad_term, log_det_term
+
+    def repeat(self, *sizes):
+        if len(sizes) < 3 or tuple(sizes[-2:]) != (1, 1):
+            raise RuntimeError(
+                "Invalid repeat arguments {}. Currently, repeat only works to create repeated "
+                "batches of a 2D LazyTensor.".format(tuple(sizes))
+            )
+
+        padded_batch_repeat = tuple(1 for _ in len(sizes) - 2 - len(self.batch_repeat)) + self.batch_repeat
+        return self.__class__(self, batch_repeat=torch.Size(
+            orig_repeat_size * new_repeat_size
+            for orig_repeat_size, new_repeat_size in zip(padded_batch_repeat, sizes[:-2]))
+        )
+
+    def __getitem__(self, index):
+        """
+        Supports subindexing of the matrix this LazyTensor represents. This may return either another
+        :obj:`gpytorch.lazy.LazyTensor` or a :obj:`torch.tensor` depending on the exact implementation.
+        """
+        args = []
+        kwargs = self.base_lazy_tensor._kwargs
+        num_base_batch_dims = len(self.base_lazy_tensor.batch_shape)
+
+        for arg in self.base_lazy_tensor._args:
+            if torch.is_tensor(arg):
+                arg_base_shape_len = max(arg.dim() - num_base_batch_dims, 0)
+                args.append(arg.repeat(*self.batch_repeat, *[1 for _ in range(arg_base_shape_len)]))
+            elif isinstance(arg, LazyTensor):
+                args.append(BatchRepeatLazyTensor(arg, batch_repeat=self.batch_repeat))
+            else:
+                args.append(arg)
+
+        new_lazy_tensor = self.base_lazy_tensor.__class__(*args, **kwargs)
+        return new_lazy_tensor.__getitem__(index)
diff --git a/gpytorch/lazy/chol_lazy_tensor.py b/gpytorch/lazy/chol_lazy_tensor.py
index 5d8f34af..f9c5b6f2 100644
--- a/gpytorch/lazy/chol_lazy_tensor.py
+++ b/gpytorch/lazy/chol_lazy_tensor.py
@@ -56,13 +56,14 @@ class CholLazyTensor(RootLazyTensor):
             res = super(CholLazyTensor, self).inv_matmul(rhs)
         return res
 
-    def inv_quad_log_det(self, inv_quad_rhs=None, log_det=False):
+    def inv_quad_log_det(self, inv_quad_rhs=None, log_det=False, reduce_inv_quad=True):
         inv_quad_term = None
         log_det_term = None
-        is_batch = self.ndimension() == 3
 
         if inv_quad_rhs is not None:
-            inv_quad_term = self.inv_matmul(inv_quad_rhs).mul(inv_quad_rhs).sum(-1).sum(-1, keepdim=(not is_batch))
+            inv_quad_term = self.inv_matmul(inv_quad_rhs).mul(inv_quad_rhs).sum(-2)
+            if reduce_inv_quad:
+                inv_quad_term = inv_quad_term.sum(-1)
 
         if log_det:
             log_det_term = self._chol_diag.log().sum(-1).mul(2)
diff --git a/gpytorch/lazy/constant_mul_lazy_tensor.py b/gpytorch/lazy/constant_mul_lazy_tensor.py
index c07906be..70e772fc 100644
--- a/gpytorch/lazy/constant_mul_lazy_tensor.py
+++ b/gpytorch/lazy/constant_mul_lazy_tensor.py
@@ -142,9 +142,6 @@ class ConstantMulLazyTensor(LazyTensor):
         res = res * self._constant_as(res)
         return res
 
-    def repeat(self, *sizes):
-        return ConstantMulLazyTensor(self.base_lazy_tensor.repeat(*sizes), self.constant)
-
     def __getitem__(self, i):
         constant = self.constant
         if constant.numel() > 1:
diff --git a/gpytorch/lazy/diag_lazy_tensor.py b/gpytorch/lazy/diag_lazy_tensor.py
index a1199a53..ac843d72 100644
--- a/gpytorch/lazy/diag_lazy_tensor.py
+++ b/gpytorch/lazy/diag_lazy_tensor.py
@@ -64,14 +64,6 @@ class DiagLazyTensor(LazyTensor):
         else:
             return super(DiagLazyTensor, self).evaluate()
 
-    def repeat(self, *sizes):
-        """
-        Repeat elements of the Tensor.
-        Right now it only works to create a batched version.
-        """
-
-        return self.__class__(self._diag.repeat(sizes[0], 1))
-
     def zero_mean_mvn_samples(self, num_samples):
         if self.ndimension() == 3:
             base_samples = torch.randn(
diff --git a/gpytorch/lazy/interpolated_lazy_tensor.py b/gpytorch/lazy/interpolated_lazy_tensor.py
index 00b6856f..679e00cf 100644
--- a/gpytorch/lazy/interpolated_lazy_tensor.py
+++ b/gpytorch/lazy/interpolated_lazy_tensor.py
@@ -561,25 +561,6 @@ class InterpolatedLazyTensor(LazyTensor):
         else:
             return super(InterpolatedLazyTensor, self).mul(other)
 
-    def repeat(self, *sizes):
-        """
-        Repeat elements of the Tensor.
-        Right now it only works to create a batched version of a InterpolatedLazyTensor.
-
-        e.g. `var.repeat(3, 1, 1)` creates a batched version of length 3
-        """
-        if not len(sizes) == 3 and sizes[1] == 1 and sizes[2] == 1:
-            raise RuntimeError("Repeat only works to create a batched version at the moment.")
-
-        return self.__class__(
-            self.base_lazy_tensor,
-            self.left_interp_indices.repeat(*sizes),
-            self.left_interp_values.repeat(*sizes),
-            self.right_interp_indices.repeat(*sizes),
-            self.right_interp_values.repeat(*sizes),
-            **self._kwargs
-        )
-
     def sum_batch(self, sum_batch_size=None):
         left_interp_indices = self.left_interp_indices
         left_interp_values = self.left_interp_values
diff --git a/gpytorch/lazy/kronecker_product_lazy_tensor.py b/gpytorch/lazy/kronecker_product_lazy_tensor.py
index cdabbabb..6ff562f3 100644
--- a/gpytorch/lazy/kronecker_product_lazy_tensor.py
+++ b/gpytorch/lazy/kronecker_product_lazy_tensor.py
@@ -224,6 +224,3 @@ class KroneckerProductLazyTensor(LazyTensor):
             left_indices = left_indices - (left_indices_i * left_size)
             right_indices = right_indices - (right_indices_i * right_size)
         return res
-
-    def repeat(self, *sizes):
-        return KroneckerProductLazyTensor(*[lazy_tensor.repeat(*sizes) for lazy_tensor in self.lazy_tensors])
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 51116e1e..1312cff2 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -575,7 +575,7 @@ class LazyTensor(object):
         res, _ = self.inv_quad_log_det(inv_quad_rhs=tensor, log_det=False)
         return res
 
-    def inv_quad_log_det(self, inv_quad_rhs=None, log_det=False):
+    def inv_quad_log_det(self, inv_quad_rhs=None, log_det=False, reduce_inv_quad=True):
         """
         Computes an inverse quadratic form (w.r.t self) with several right hand sides.
         I.e. computes tr( tensor^T self^{-1} tensor )
@@ -618,7 +618,7 @@ class LazyTensor(object):
         if inv_quad_rhs is not None:
             args = [inv_quad_rhs] + list(args)
 
-        res = InvQuadLogDet(
+        inv_quad_term, log_det_term = InvQuadLogDet(
             representation_tree=self.representation_tree(),
             matrix_shape=self.matrix_shape,
             batch_shape=self.batch_shape,
@@ -629,7 +629,10 @@ class LazyTensor(object):
             preconditioner=self._preconditioner()[0],
             log_det_correction=self._preconditioner()[1],
         )(*args)
-        return res
+
+        if inv_quad_term.numel() and reduce_inv_quad:
+            inv_quad_term = inv_quad_term.sum(-1)
+        return inv_quad_term, log_det_term
 
     @property
     def is_square(self):
@@ -810,6 +813,32 @@ class LazyTensor(object):
         """
         return self.shape.numel()
 
+    def repeat(self, *sizes):
+        """
+        Repeats this tensor along the specified dimensions.
+
+        Currently, this only works to create repeated batches of a 2D LazyTensor.
+        I.e. all calls should be `lazy_tensor.repeat(<size>, 1, 1)`.
+
+        Example:
+            >>> lazy_tensor = gpytorch.lazy.ToeplitzLazyTensor(torch.tensor([4. 1., 0.5]))
+            >>> lazy_tensor.repeat(2, 1, 1).evaluate()
+            tensor([[[4.0000, 1.0000, 0.5000],
+                     [1.0000, 4.0000, 1.0000],
+                     [0.5000, 1.0000, 4.0000]],
+                    [[4.0000, 1.0000, 0.5000],
+                     [1.0000, 4.0000, 1.0000],
+                     [0.5000, 1.0000, 4.0000]]])
+        """
+        if len(sizes) < 3 or tuple(sizes[-2:]) != (1, 1):
+            raise RuntimeError(
+                "Invalid repeat arguments {}. Currently, repeat only works to create repeated "
+                "batches of a 2D LazyTensor.".format(tuple(sizes))
+            )
+
+        from .batch_repeat_lazy_tensor import BatchRepeatLazyTensor
+        return BatchRepeatLazyTensor(self, batch_repeat=torch.Size(sizes[:-2]))
+
     def representation(self):
         """
         Returns the Tensors that are used to define the LazyTensor
@@ -1227,7 +1256,7 @@ class LazyTensor(object):
         if torch.is_tensor(val) or isinstance(val, LazyTensor):
             if not hasattr(self, "_args"):
                 raise RuntimeError(
-                    "Cannot assign {name} to LazyTensor before calling " "LazyTensor.__init__()".format(name=name)
+                    "Cannot assign {name} to LazyTensor before calling LazyTensor.__init__()".format(name=name)
                 )
         object.__setattr__(self, name, val)
 
diff --git a/gpytorch/lazy/non_lazy_tensor.py b/gpytorch/lazy/non_lazy_tensor.py
index 4eb9f646..c29ce902 100644
--- a/gpytorch/lazy/non_lazy_tensor.py
+++ b/gpytorch/lazy/non_lazy_tensor.py
@@ -57,9 +57,6 @@ class NonLazyTensor(LazyTensor):
     def evaluate(self):
         return self.tensor
 
-    def repeat(self, *sizes):
-        return NonLazyTensor(self.tensor.repeat(*sizes))
-
     def __getitem__(self, index):
         res = self.tensor[index]
         if not isinstance(index, tuple):
diff --git a/gpytorch/lazy/root_lazy_tensor.py b/gpytorch/lazy/root_lazy_tensor.py
index fc2fd433..fb2621f2 100644
--- a/gpytorch/lazy/root_lazy_tensor.py
+++ b/gpytorch/lazy/root_lazy_tensor.py
@@ -109,14 +109,6 @@ class RootLazyTensor(LazyTensor):
             self._evaluated_memo = torch.matmul(self._evaluated_root, self._evaluated_root.transpose(-1, -2))
         return self._evaluated_memo
 
-    def repeat(self, *sizes):
-        """
-        Repeat elements of the Tensor.
-        Right now it only works to create a batched version.
-        """
-
-        return self.__class__(self.root.repeat(*sizes))
-
     def root_decomposition_size(self):
         return self.root.size(-1)
 
diff --git a/gpytorch/lazy/sum_lazy_tensor.py b/gpytorch/lazy/sum_lazy_tensor.py
index 39a5159e..3fe02afc 100644
--- a/gpytorch/lazy/sum_lazy_tensor.py
+++ b/gpytorch/lazy/sum_lazy_tensor.py
@@ -49,9 +49,9 @@ class SumLazyTensor(LazyTensor):
     def _get_indices(self, left_indices, right_indices):
         return sum(lazy_tensor._get_indices(left_indices, right_indices) for lazy_tensor in self.lazy_tensors)
 
-    def add_jitter(self):
+    def add_jitter(self, jitter_val=1e-3):
         lazy_tensors = list(self.lazy_tensors[:-1])
-        lazy_tensors.append(self.lazy_tensors[-1].add_jitter())
+        lazy_tensors.append(self.lazy_tensors[-1].add_jitter(jitter_val=jitter_val))
         return SumLazyTensor(*lazy_tensors)
 
     def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root, test_train_covar):
@@ -92,9 +92,6 @@ class SumLazyTensor(LazyTensor):
         res = res.view(size)
         return res
 
-    def repeat(self, *sizes):
-        return self.__class__(*(lazy_tensor.repeat(*sizes) for lazy_tensor in self.lazy_tensors))
-
     def sum_batch(self, sum_batch_size=None):
         return self.__class__(*(lazy_tensor.sum_batch(sum_batch_size) for lazy_tensor in self.lazy_tensors))
 
diff --git a/gpytorch/lazy/toeplitz_lazy_tensor.py b/gpytorch/lazy/toeplitz_lazy_tensor.py
index 935ff677..6af7ac8e 100644
--- a/gpytorch/lazy/toeplitz_lazy_tensor.py
+++ b/gpytorch/lazy/toeplitz_lazy_tensor.py
@@ -49,9 +49,9 @@ class ToeplitzLazyTensor(LazyTensor):
         toeplitz_indices = (left_indices - right_indices).fmod(n_grid).abs().long()
         return self.column.index_select(0, toeplitz_indices)
 
-    def add_jitter(self):
+    def add_jitter(self, jitter_val=1e-3):
         jitter = torch.zeros_like(self.column)
-        jitter.narrow(-1, 0, 1).fill_(1e-4)
+        jitter.narrow(-1, 0, 1).fill_(jitter_val)
         return ToeplitzLazyTensor(self.column.add(jitter))
 
     def diag(self):
@@ -62,13 +62,3 @@ class ToeplitzLazyTensor(LazyTensor):
         if self.column.ndimension() > 1:
             diag_term = diag_term.unsqueeze(-1)
         return diag_term.expand(*self.column.size())
-
-    def repeat(self, *sizes):
-        """
-        Repeat elements of the Tensor.
-        Right now it only works to create a batched version of a ToeplitzLazyTensor.
-
-        e.g. `var.repeat(3, 1, 1)` creates a batched version of length 3
-        """
-
-        return ToeplitzLazyTensor(self.column.repeat(sizes[0], 1))
diff --git a/gpytorch/lazy/zero_lazy_tensor.py b/gpytorch/lazy/zero_lazy_tensor.py
index 5c2a43a0..00d71be7 100644
--- a/gpytorch/lazy/zero_lazy_tensor.py
+++ b/gpytorch/lazy/zero_lazy_tensor.py
@@ -97,7 +97,7 @@ class ZeroLazyTensor(LazyTensor):
     def inv_quad(self, tensor):
         raise RuntimeError("ZeroLazyTensors are not invertible!")
 
-    def inv_quad_log_det(self, inv_quad_rhs=None, log_det=False):
+    def inv_quad_log_det(self, inv_quad_rhs=None, log_det=False, reduce_inv_quad=True):
         raise RuntimeError("ZeroLazyTensors are not invertible!")
 
     def log_det(self):


Test patch:
diff --git a/test/lazy/_lazy_tensor_test_case.py b/test/lazy/_lazy_tensor_test_case.py
index a50c1f01..b0820668 100644
--- a/test/lazy/_lazy_tensor_test_case.py
+++ b/test/lazy/_lazy_tensor_test_case.py
@@ -199,6 +199,26 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
         diff = (res - actual).abs() / actual.abs().clamp(1, 1e10)
         self.assertLess(diff.item(), 15e-2)
 
+    def test_inv_quad_log_det_no_reduce(self):
+        # Forward
+        lazy_tensor = self.create_lazy_tensor()
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor)
+
+        vecs = torch.randn(lazy_tensor.size(1), 3, requires_grad=True)
+        vecs_copy = vecs.clone()
+
+        with gpytorch.settings.num_trace_samples(128):
+            res_inv_quad, res_log_det = lazy_tensor.inv_quad_log_det(
+                inv_quad_rhs=vecs, log_det=True, reduce_inv_quad=False
+            )
+        res = res_inv_quad.sum(-1) + res_log_det
+
+        actual_inv_quad = evaluated.inverse().matmul(vecs_copy).mul(vecs_copy).sum()
+        actual = actual_inv_quad + torch.logdet(evaluated)
+
+        diff = (res - actual).abs() / actual.abs().clamp(1, 1e10)
+        self.assertLess(diff.item(), 15e-2)
+
     def test_sample(self):
         if self.__class__.should_test_sample:
             lazy_tensor = self.create_lazy_tensor()
@@ -405,6 +425,33 @@ class BatchLazyTensorTestCase(RectangularBatchLazyTensorTestCase):
         for i in range(diffs.numel()):
             self.assertLess(diffs[i].item(), 15e-2)
 
+    def test_inv_quad_log_det_no_reduce(self):
+        # Forward
+        lazy_tensor = self.create_lazy_tensor()
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor)
+
+        vecs = torch.randn(lazy_tensor.size(0), lazy_tensor.size(1), 3, requires_grad=True)
+        vecs_copy = vecs.clone()
+
+        with gpytorch.settings.num_trace_samples(128):
+            res_inv_quad, res_log_det = lazy_tensor.inv_quad_log_det(
+                inv_quad_rhs=vecs, log_det=True, reduce_inv_quad=False
+            )
+        res = res_inv_quad.sum(-1) + res_log_det
+
+        actual_inv_quad = torch.cat(
+            [
+                evaluated[i].inverse().matmul(vecs_copy[i]).mul(vecs_copy[i]).sum().unsqueeze(0)
+                for i in range(lazy_tensor.size(0))
+            ]
+        )
+        actual = actual_inv_quad + torch.cat(
+            [torch.logdet(evaluated[i]).unsqueeze(0) for i in range(lazy_tensor.size(0))]
+        )
+        diffs = (res - actual).abs() / actual.abs().clamp(1, 1e10)
+        for i in range(lazy_tensor.size(0)):
+            self.assertLess(diffs[i].item(), 15e-2)
+
     def test_sample(self):
         if self.__class__.should_test_sample:
             lazy_tensor = self.create_lazy_tensor()
diff --git a/test/lazy/test_batch_repeat_lazy_tensor.py b/test/lazy/test_batch_repeat_lazy_tensor.py
new file mode 100644
index 00000000..b80cf739
--- /dev/null
+++ b/test/lazy/test_batch_repeat_lazy_tensor.py
@@ -0,0 +1,41 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import torch
+import unittest
+from gpytorch.lazy import ToeplitzLazyTensor, BatchRepeatLazyTensor
+from test.lazy._lazy_tensor_test_case import BatchLazyTensorTestCase
+
+
+class TestBatchRepeatLazyTensor(BatchLazyTensorTestCase, unittest.TestCase):
+    seed = 0
+
+    def create_lazy_tensor(self):
+        toeplitz_column = torch.tensor(
+            [4, 0, 0, 1], dtype=torch.float, requires_grad=True
+        )
+        return BatchRepeatLazyTensor(ToeplitzLazyTensor(toeplitz_column), torch.Size((3,)))
+
+    def evaluate_lazy_tensor(self, lazy_tensor):
+        evaluated = lazy_tensor.base_lazy_tensor.evaluate()
+        return evaluated.repeat(*lazy_tensor.batch_repeat, 1, 1)
+
+
+class TestBatchRepeatLazyTensorBatch(BatchLazyTensorTestCase, unittest.TestCase):
+    seed = 0
+
+    def create_lazy_tensor(self):
+        toeplitz_column = torch.tensor(
+            [[4, 0, 0, 1], [3, 0, -0.5, -1]], dtype=torch.float, requires_grad=True
+        )
+        return BatchRepeatLazyTensor(ToeplitzLazyTensor(toeplitz_column), torch.Size((3,)))
+
+    def evaluate_lazy_tensor(self, lazy_tensor):
+        evaluated = lazy_tensor.base_lazy_tensor.evaluate()
+        return evaluated.repeat(*lazy_tensor.batch_repeat, 1, 1)
+
+
+if __name__ == "__main__":
+    unittest.main()