Repository: cornellius-gp/gpytorch
Pull request #459 (2019-01-09 21:08:37)

 LazyTensors can broadcast matmuls
This applies to all LazyTensors except: `BatchRepeatLazyTensor`, `BlockDiagonalLazyTensor`,`SumBatchLazyTensor`.


Patch:
diff --git a/gpytorch/functions/_dsmm.py b/gpytorch/functions/_dsmm.py
index 765ac21f..b9af8deb 100644
--- a/gpytorch/functions/_dsmm.py
+++ b/gpytorch/functions/_dsmm.py
@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 
-import torch
 from torch.autograd import Function
 from ..utils.sparse import bdsmm
 
@@ -10,13 +9,7 @@ class DSMM(Function):
         self.sparse = sparse
 
     def forward(self, dense):
-        if self.sparse.ndimension() == 3:
-            return bdsmm(self.sparse, dense)
-        else:
-            return torch.dsmm(self.sparse, dense)
+        return bdsmm(self.sparse, dense)
 
     def backward(self, grad_output):
-        if self.sparse.ndimension() == 3:
-            return bdsmm(self.sparse.transpose(1, 2), grad_output)
-        else:
-            return torch.dsmm(self.sparse.t(), grad_output)
+        return bdsmm(self.sparse.transpose(-1, -2), grad_output)
diff --git a/gpytorch/functions/_matmul.py b/gpytorch/functions/_matmul.py
index 9a170832..9e00d8f9 100644
--- a/gpytorch/functions/_matmul.py
+++ b/gpytorch/functions/_matmul.py
@@ -57,6 +57,8 @@ class Matmul(Function):
             else:
                 rhs_grad = lazy_tsr._t_matmul(grad_output)
 
-            rhs_grad = rhs_grad.view(rhs_shape)
+            # For broadcasting
+            if rhs_grad.dim() > len(rhs_shape):
+                rhs_grad = rhs_grad.contiguous().view(-1, *rhs_shape).sum(0)
 
         return tuple([rhs_grad] + list(arg_grads))
diff --git a/gpytorch/lazy/constant_mul_lazy_tensor.py b/gpytorch/lazy/constant_mul_lazy_tensor.py
index 16029dcb..3407810b 100644
--- a/gpytorch/lazy/constant_mul_lazy_tensor.py
+++ b/gpytorch/lazy/constant_mul_lazy_tensor.py
@@ -65,13 +65,11 @@ class ConstantMulLazyTensor(LazyTensor):
 
     def _approx_diag(self):
         res = self.base_lazy_tensor._approx_diag()
-        res_mat_shape = res.shape[len(self.base_lazy_tensor.batch_shape) :]
-        constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
-        return res * constant
+        return res * self._constant.unsqueeze(-1)
 
     def _get_indices(self, left_indices, right_indices, *batch_indices):
         res = self.base_lazy_tensor._get_indices(left_indices, right_indices, *batch_indices)
-        constant = self.constant.__getitem__(batch_indices)
+        constant = self._constant.expand(self.batch_shape).__getitem__(batch_indices)
         return res * constant
 
     def _getitem(self, *indices):
@@ -81,7 +79,7 @@ class ConstantMulLazyTensor(LazyTensor):
         # This effects runntimes by up to 5x on simple exat GPs
         # Run __getitem__ on the base_lazy_tensor and the constant
         base_lazy_tensor = self.base_lazy_tensor._getitem(*indices)
-        constant = self.constant[indices[:-2]]
+        constant = self._constant.expand(self.batch_shape)[indices[:-2]]
 
         if torch.is_tensor(base_lazy_tensor):
             constant = constant.view(*constant.shape, *[1] * (base_lazy_tensor.dim() - constant.dim()))
@@ -90,9 +88,7 @@ class ConstantMulLazyTensor(LazyTensor):
 
     def _matmul(self, rhs):
         res = self.base_lazy_tensor._matmul(rhs)
-        res_mat_shape = res.shape[len(self.base_lazy_tensor.batch_shape) :]
-        constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
-        res = res * constant
+        res = res * self.expanded_constant
         return res
 
     def _quad_form_derivative(self, left_vecs, right_vecs):
@@ -101,9 +97,7 @@ class ConstantMulLazyTensor(LazyTensor):
         constant_deriv = constant_deriv.sum(-2).sum(-1)
 
         # Get derivaties of everything else
-        res_mat_shape = left_vecs.shape[len(self.base_lazy_tensor.batch_shape) :]
-        constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
-        left_vecs = left_vecs * constant
+        left_vecs = left_vecs * self.expanded_constant
         res = self.base_lazy_tensor._quad_form_derivative(left_vecs, right_vecs)
 
         return res + (constant_deriv,)
@@ -113,19 +107,17 @@ class ConstantMulLazyTensor(LazyTensor):
 
     def _t_matmul(self, rhs):
         res = self.base_lazy_tensor._t_matmul(rhs)
-        res_mat_shape = res.shape[len(self.base_lazy_tensor.batch_shape) :]
-        constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
-        res = res * constant
+        res = res * self.expanded_constant
         return res
 
     def _transpose_nonbatch(self):
-        return ConstantMulLazyTensor(self.base_lazy_tensor._transpose_nonbatch(), self.constant)
+        return ConstantMulLazyTensor(self.base_lazy_tensor._transpose_nonbatch(), self._constant)
 
     @property
-    def constant(self):
+    def expanded_constant(self):
         # Make sure that the constant can be expanded to the appropriate size
         try:
-            constant = self._constant.expand(self.base_lazy_tensor.batch_shape)
+            constant = self._constant.view(*self._constant.shape, 1, 1)
         except RuntimeError:
             raise RuntimeError(
                 "ConstantMulLazyTensor of size {} received an invalid constant of size {}.".format(
@@ -137,13 +129,9 @@ class ConstantMulLazyTensor(LazyTensor):
 
     def diag(self):
         res = self.base_lazy_tensor.diag()
-        res_mat_shape = res.shape[len(self.base_lazy_tensor.batch_shape) :]
-        constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
-        return res * constant
+        return res * self._constant.unsqueeze(-1)
 
     @cached
     def evaluate(self):
         res = self.base_lazy_tensor.evaluate()
-        res_mat_shape = res.shape[len(self.base_lazy_tensor.batch_shape) :]
-        constant = self.constant.view(*self.constant.shape, *[1 for i in res_mat_shape])
-        return constant * res
+        return res * self.expanded_constant
diff --git a/gpytorch/lazy/interpolated_lazy_tensor.py b/gpytorch/lazy/interpolated_lazy_tensor.py
index 59e1e91b..b37e08c9 100644
--- a/gpytorch/lazy/interpolated_lazy_tensor.py
+++ b/gpytorch/lazy/interpolated_lazy_tensor.py
@@ -6,6 +6,7 @@ from .lazy_tensor import LazyTensor
 from .non_lazy_tensor import lazify
 from .root_lazy_tensor import RootLazyTensor
 from ..utils import sparse
+from ..utils.batch import _create_batch_indices
 from ..utils.interpolation import left_interp, left_t_interp
 from ..utils.sparse import bdsmm
 
@@ -22,19 +23,10 @@ class InterpolatedLazyTensor(LazyTensor):
         base_lazy_tensor = lazify(base_lazy_tensor)
 
         if left_interp_indices is None:
-            num_rows = base_lazy_tensor.size()[-2]
+            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)
-            if base_lazy_tensor.ndimension() == 3:
-                left_interp_indices = left_interp_indices.unsqueeze(0).expand(base_lazy_tensor.size(0), num_rows, 1)
-            elif right_interp_indices is not None and right_interp_indices.ndimension() == 3:
-                left_interp_indices = left_interp_indices.unsqueeze(0).expand(right_interp_indices.size(0), num_rows, 1)
-        else:
-            if left_interp_indices.dim() != base_lazy_tensor.ndimension():
-                raise RuntimeError(
-                    "Expected left_interp_indices ({}) to have the same number of dimensions as "
-                    "base_lazy_Tensor ({})".format(left_interp_indices.size(), base_lazy_tensor.size())
-                )
+            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(
@@ -48,21 +40,10 @@ class InterpolatedLazyTensor(LazyTensor):
                 )
 
         if right_interp_indices is None:
-            num_rows = base_lazy_tensor.size()[-2]
+            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)
-            if base_lazy_tensor.ndimension() == 3:
-                right_interp_indices = right_interp_indices.unsqueeze(0).expand(base_lazy_tensor.size(0), num_rows, 1)
-            elif left_interp_indices.ndimension() == 3:
-                right_interp_indices = right_interp_indices.unsqueeze(0).expand(
-                    left_interp_indices.size(0), num_rows, 1
-                )
-        else:
-            if left_interp_indices.dim() != base_lazy_tensor.ndimension():
-                raise RuntimeError(
-                    "Expected left_interp_indices ({}) to have the same number of dimensions as "
-                    "base_lazy_Tensor ({})".format(left_interp_indices.size(), base_lazy_tensor.size())
-                )
+            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(
@@ -75,6 +56,18 @@ class InterpolatedLazyTensor(LazyTensor):
                     "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
         )
@@ -164,39 +157,38 @@ class InterpolatedLazyTensor(LazyTensor):
         n_left_interp = self.left_interp_indices.size(-1)
         n_right_interp = self.right_interp_indices.size(-1)
         n_inducing = right_res.size(-2)
-        if self.left_interp_indices.ndimension() == 3:
-            batch_size = self.left_interp_indices.size(0)
 
         # left_interp_values grad
         right_interp_right_res = self.base_lazy_tensor._matmul(right_res).contiguous()
-        if self.left_interp_indices.ndimension() == 3:
-            batch_offset = torch.arange(0, batch_size, dtype=torch.long, device=self.device)
+        batch_shape = torch.Size(right_interp_right_res.shape[:-2])
+        batch_size = batch_shape.numel()
+        if len(batch_shape):
+            batch_offset = torch.arange(0, batch_size, dtype=torch.long, device=self.device).view(*batch_shape)
             batch_offset.unsqueeze_(-1).unsqueeze_(-1).mul_(n_inducing)
-
+            batched_right_interp_indices = self.right_interp_indices
             batched_left_interp_indices = (self.left_interp_indices + batch_offset).view(-1)
-            flattened_right_interp_right_res = right_interp_right_res.view(batch_size * n_inducing, n_vecs)
-
-            selected_right_vals = flattened_right_interp_right_res.index_select(0, batched_left_interp_indices)
-            selected_right_vals = selected_right_vals.view(batch_size, n_left_rows, n_left_interp, n_vecs)
         else:
-            selected_right_vals = right_interp_right_res.index_select(0, self.left_interp_indices.view(-1))
-            selected_right_vals = selected_right_vals.view(n_left_rows, n_left_interp, n_vecs)
+            batched_left_interp_indices = self.left_interp_indices.view(-1)
+
+        flattened_right_interp_right_res = right_interp_right_res.view(batch_size * n_inducing, n_vecs)
+        selected_right_vals = flattened_right_interp_right_res.index_select(0, batched_left_interp_indices)
+        selected_right_vals = selected_right_vals.view(*batch_shape, n_left_rows, n_left_interp, n_vecs)
         left_values_grad = (selected_right_vals * left_vecs.unsqueeze(-2)).sum(-1)
 
         # right_interp_values_grad
         left_interp_left_res = self.base_lazy_tensor._t_matmul(left_res).contiguous()
-        if self.right_interp_indices.ndimension() == 3:
-            batch_offset = torch.arange(0, batch_size, dtype=torch.long, device=self.device)
+        batch_shape = left_interp_left_res.shape[:-2]
+        batch_size = batch_shape.numel()
+        if len(batch_shape):
+            batch_offset = torch.arange(0, batch_size, dtype=torch.long, device=self.device).view(*batch_shape)
             batch_offset.unsqueeze_(-1).unsqueeze_(-1).mul_(n_inducing)
-
             batched_right_interp_indices = (self.right_interp_indices + batch_offset).view(-1)
-            flattened_left_interp_left_res = left_interp_left_res.view(batch_size * n_inducing, n_vecs)
-
-            selected_left_vals = flattened_left_interp_left_res.index_select(0, batched_right_interp_indices)
-            selected_left_vals = selected_left_vals.view(batch_size, n_right_rows, n_right_interp, n_vecs)
         else:
-            selected_left_vals = left_interp_left_res.index_select(0, self.right_interp_indices.view(-1))
-            selected_left_vals = selected_left_vals.view(n_right_rows, n_right_interp, n_vecs)
+            batched_right_interp_indices = self.right_interp_indices.view(-1)
+
+        flattened_left_interp_left_res = left_interp_left_res.view(batch_size * n_inducing, n_vecs)
+        selected_left_vals = flattened_left_interp_left_res.index_select(0, batched_right_interp_indices)
+        selected_left_vals = selected_left_vals.view(*batch_shape, n_right_rows, n_right_interp, n_vecs)
         right_values_grad = (selected_left_vals * right_vecs.unsqueeze(-2)).sum(-1)
 
         # Return zero grad for interp indices
@@ -212,12 +204,9 @@ class InterpolatedLazyTensor(LazyTensor):
         return res
 
     def _size(self):
-        if self.left_interp_indices.ndimension() == 3:
-            return torch.Size(
-                (self.left_interp_indices.size(0), self.left_interp_indices.size(1), self.right_interp_indices.size(1))
-            )
-        else:
-            return torch.Size((self.left_interp_indices.size(0), self.right_interp_indices.size(0)))
+        return torch.Size(
+            self.base_lazy_tensor.batch_shape + (self.left_interp_indices.size(-2), self.right_interp_indices.size(-2))
+        )
 
     def _transpose_nonbatch(self):
         res = self.__class__(
@@ -285,57 +274,38 @@ class InterpolatedLazyTensor(LazyTensor):
         return self._sparse_right_interp_t_memo
 
     def diag(self):
+        if self.left_interp_indices.size(-2) != self.right_interp_indices.size(-2):
+            raise RuntimeError(
+                "diag is only defined for square LazyTensors - won't work for LazyTensor of size {}".format(self.shape)
+            )
+
         if isinstance(self.base_lazy_tensor, RootLazyTensor):
             res = left_interp(self.left_interp_indices, self.left_interp_values, self.base_lazy_tensor.root.evaluate())
             return res.pow(2).sum(-1)
         else:
-            batch_size = None
-            n_data = None
-            n_interp = None
-            if self.left_interp_indices.ndimension() == 3:
-                batch_size, n_data, n_interp = self.left_interp_indices.size()
-            else:
-                n_data, n_interp = self.left_interp_indices.size()
+            num_data, num_left_interp = self.left_interp_indices.shape[-2:]
+            num_right_interp = self.left_interp_indices.size(-1)
 
             # Batch compute the non-zero values of the outer products w_left^k w_right^k^T
             left_interp_values = self.left_interp_values.unsqueeze(-1)
             right_interp_values = self.right_interp_values.unsqueeze(-2)
             interp_values = torch.matmul(left_interp_values, right_interp_values).view(-1)
-
-            # Batch compute indicies that will be non-zero for row k
-            if batch_size is None:
-                left_interp_indices = self.left_interp_indices.unsqueeze(-1).expand(n_data, n_interp, n_interp)
-                right_interp_indices = self.right_interp_indices.unsqueeze(-2).expand(n_data, n_interp, n_interp)
-            else:
-                left_interp_indices = self.left_interp_indices.unsqueeze(-1).expand(
-                    batch_size, n_data, n_interp, n_interp
-                )
-                right_interp_indices = self.right_interp_indices.unsqueeze(-2).expand(
-                    batch_size, n_data, n_interp, n_interp
-                )
-
-            left_interp_indices = left_interp_indices.contiguous().view(-1)
-            right_interp_indices = right_interp_indices.contiguous().view(-1)
-
-            base_batch_indices = (
-                [
-                    torch.arange(0, batch_size, dtype=torch.long, device=self.device)
-                    .unsqueeze(-1)
-                    .repeat(1, n_data * n_interp * n_interp)
-                    .view(-1)
-                ]
-                if batch_size is not None
-                else []
+            left_interp_indices = self.left_interp_indices.unsqueeze(-1).expand(
+                *self.batch_shape, num_data, num_left_interp, num_right_interp
+            ).contiguous().view(-1)
+            right_interp_indices = self.right_interp_indices.unsqueeze(-2).expand(
+                *self.batch_shape, num_data, num_left_interp, num_right_interp
+            ).contiguous().view(-1)
+
+            batch_indices = _create_batch_indices(
+                self.batch_shape, num_data * num_left_interp * num_right_interp, self.device
             )
             base_var_vals = self.base_lazy_tensor._get_indices(
-                left_interp_indices, right_interp_indices, *base_batch_indices
+                left_interp_indices, right_interp_indices, *batch_indices
             )
             res = interp_values * base_var_vals
 
-            if batch_size is None:
-                res = res.view(n_data, -1).sum(-1)
-            else:
-                res = res.view(batch_size, n_data, -1).sum(-1)
+            res = res.view(*self.batch_shape, num_data, -1).sum(-1)
             return res
 
     def matmul(self, tensor):
@@ -438,16 +408,11 @@ class InterpolatedLazyTensor(LazyTensor):
 
     def zero_mean_mvn_samples(self, num_samples):
         base_samples = self.base_lazy_tensor.zero_mean_mvn_samples(num_samples)
-        if self.ndimension() == 3:
-            res = left_interp(
-                self.left_interp_indices, self.left_interp_values, base_samples.permute(1, 2, 0).contiguous()
-            )
-            return res.permute(2, 0, 1).contiguous()
-        else:
-            res = left_interp(
-                self.left_interp_indices, self.left_interp_values, base_samples.permute(1, 0).contiguous()
-            )
-            return res.permute(1, 0).contiguous()
+        batch_iter = tuple(range(1, base_samples.dim()))
+        base_samples = base_samples.permute(*batch_iter, 0)
+        res = left_interp(self.left_interp_indices, self.left_interp_values, base_samples).contiguous()
+        batch_iter = tuple(range(res.dim() - 1))
+        return res.permute(-1, *batch_iter).contiguous()
 
     def _getitem_nonbatch(self, row_index, col_index, first_tensor_index_dim=None):
         """
diff --git a/gpytorch/lazy/kronecker_product_lazy_tensor.py b/gpytorch/lazy/kronecker_product_lazy_tensor.py
index 3455ce7e..4ed1a581 100644
--- a/gpytorch/lazy/kronecker_product_lazy_tensor.py
+++ b/gpytorch/lazy/kronecker_product_lazy_tensor.py
@@ -4,6 +4,8 @@ import torch
 import operator
 from .lazy_tensor import LazyTensor
 from .non_lazy_tensor import lazify
+from ..utils.broadcasting import _matmul_broadcast_shape
+from ..utils.memoize import cached
 from functools import reduce
 
 
@@ -11,25 +13,32 @@ def _prod(iterable):
     return reduce(operator.mul, iterable, 1)
 
 
-def _matmul(lazy_tensors, rhs):
-    res = rhs.contiguous()
+def _matmul(lazy_tensors, kp_shape, rhs):
+    output_shape = _matmul_broadcast_shape(kp_shape, rhs.shape)
+    output_batch_shape = output_shape[:-2]
+
+    res = rhs.contiguous().expand(*output_batch_shape, *rhs.shape[-2:])
     num_cols = rhs.size(-1)
     for lazy_tensor in lazy_tensors:
-        res = res.view(*lazy_tensor.batch_shape, lazy_tensor.size(-1), -1)
+        res = res.view(*output_batch_shape, lazy_tensor.size(-1), -1)
         factor = lazy_tensor._matmul(res)
-        factor = factor.view(*lazy_tensor.batch_shape, lazy_tensor.size(-2), -1, num_cols).transpose(-3, -2)
-        res = factor.contiguous().view(*lazy_tensor.batch_shape, -1, num_cols)
+        factor = factor.view(*output_batch_shape, lazy_tensor.size(-2), -1, num_cols).transpose(-3, -2)
+        res = factor.contiguous().view(*output_batch_shape, -1, num_cols)
     return res
 
 
-def _t_matmul(lazy_tensors, rhs):
-    res = rhs.contiguous()
+def _t_matmul(lazy_tensors, kp_shape, rhs):
+    kp_t_shape = (*kp_shape[:-2], kp_shape[-1], kp_shape[-2])
+    output_shape = _matmul_broadcast_shape(kp_t_shape, rhs.shape)
+    output_batch_shape = torch.Size(output_shape[:-2])
+
+    res = rhs.contiguous().expand(*output_batch_shape, *rhs.shape[-2:])
     num_cols = rhs.size(-1)
     for lazy_tensor in lazy_tensors:
-        res = res.view(*lazy_tensor.batch_shape, lazy_tensor.size(-2), -1)
+        res = res.view(*output_batch_shape, lazy_tensor.size(-2), -1)
         factor = lazy_tensor._t_matmul(res)
-        factor = factor.view(*lazy_tensor.batch_shape, lazy_tensor.size(-1), -1, num_cols).transpose(-3, -2)
-        res = factor.contiguous().view(*lazy_tensor.batch_shape, -1, num_cols)
+        factor = factor.view(*output_batch_shape, lazy_tensor.size(-1), -1, num_cols).transpose(-3, -2)
+        res = factor.contiguous().view(*output_batch_shape, -1, num_cols)
     return res
 
 
@@ -40,17 +49,11 @@ class KroneckerProductLazyTensor(LazyTensor):
         except TypeError:
             raise RuntimeError("KroneckerProductLazyTensor is intended to wrap lazy tensors.")
         for prev_lazy_tensor, curr_lazy_tensor in zip(lazy_tensors[:-1], lazy_tensors[1:]):
-            if prev_lazy_tensor.ndimension() != curr_lazy_tensor.ndimension():
+            if prev_lazy_tensor.batch_shape != curr_lazy_tensor.batch_shape:
                 raise RuntimeError(
                     "KroneckerProductLazyTensor expects lazy tensors with the "
-                    "same number of dimensions. Got {}.".format([lv.ndimension() for lv in lazy_tensors])
+                    "same batch shapes. Got {}.".format([lv.batch_shape for lv in lazy_tensors])
                 )
-            if curr_lazy_tensor.ndimension() >= 3:
-                if prev_lazy_tensor.shape[:-2] != curr_lazy_tensor.shape[:-2]:
-                    raise RuntimeError(
-                        "KroneckerProductLazyTensor expects the same batch sizes for component tensors. "
-                        "Got sizes: {}".format([lv.size() for lv in lazy_tensors])
-                    )
         super(KroneckerProductLazyTensor, self).__init__(*lazy_tensors)
         self.lazy_tensors = lazy_tensors
 
@@ -59,7 +62,7 @@ class KroneckerProductLazyTensor(LazyTensor):
         if is_vec:
             rhs = rhs.unsqueeze(-1)
 
-        res = _matmul(self.lazy_tensors, rhs.contiguous())
+        res = _matmul(self.lazy_tensors, self.shape, rhs.contiguous())
 
         if is_vec:
             res = res.squeeze(-1)
@@ -70,12 +73,13 @@ class KroneckerProductLazyTensor(LazyTensor):
         if is_vec:
             rhs = rhs.unsqueeze(-1)
 
-        res = _t_matmul(self.lazy_tensors, rhs.contiguous())
+        res = _t_matmul(self.lazy_tensors, self.shape, rhs.contiguous())
 
         if is_vec:
             res = res.squeeze(-1)
         return res
 
+    @cached(name="size")
     def _size(self):
         left_size = _prod(lazy_tensor.size(-2) for lazy_tensor in self.lazy_tensors)
         right_size = _prod(lazy_tensor.size(-1) for lazy_tensor in self.lazy_tensors)
diff --git a/gpytorch/lazy/mul_lazy_tensor.py b/gpytorch/lazy/mul_lazy_tensor.py
index 5c34765f..84b8984c 100644
--- a/gpytorch/lazy/mul_lazy_tensor.py
+++ b/gpytorch/lazy/mul_lazy_tensor.py
@@ -3,6 +3,7 @@
 import torch
 
 from ..utils import prod
+from ..utils.broadcasting import _matmul_broadcast_shape
 from ..utils.memoize import cached
 from .lazy_tensor import LazyTensor
 from .non_lazy_tensor import lazify, NonLazyTensor
@@ -106,6 +107,9 @@ class MulLazyTensor(LazyTensor):
         pass
 
     def _matmul(self, rhs):
+        output_shape = _matmul_broadcast_shape(self.shape, rhs.shape)
+        output_batch_shape = output_shape[:-2]
+
         if self.non_lazy_self is not None:
             return self.non_lazy_self._matmul(rhs)
 
@@ -113,7 +117,6 @@ class MulLazyTensor(LazyTensor):
         if rhs.ndimension() == 1:
             rhs = rhs.unsqueeze(1)
             is_vector = True
-        batch_size = max(rhs.size(0), self.size(0)) if rhs.ndimension() == 3 else None
 
         # Here we have a root decomposition
         if isinstance(self.left_lazy_tensor, RootLazyTensor):
@@ -124,9 +127,9 @@ class MulLazyTensor(LazyTensor):
             n = self.size(-1)
             m = rhs.size(-1)
             # Now implement the formula (A . B) v = diag(A D_v B)
-            left_res = left_res.view(n, rank * m) if batch_size is None else left_res.view(batch_size, n, rank * m)
+            left_res = left_res.view(*output_batch_shape, n, rank * m)
             left_res = self.right_lazy_tensor._matmul(left_res)
-            left_res = left_res.view(n, rank, m) if batch_size is None else left_res.view(batch_size, n, rank, m)
+            left_res = left_res.view(*output_batch_shape, n, rank, m)
             res = left_res.mul_(left_root.unsqueeze(-1)).sum(-2)
         # This is the case where we're not doing a root decomposition, because the matrix is too small
         else:
@@ -141,20 +144,13 @@ class MulLazyTensor(LazyTensor):
         if left_vecs.ndimension() == 1:
             left_vecs = left_vecs.unsqueeze(1)
             right_vecs = right_vecs.unsqueeze(1)
-        batch_size = self.size(0) if self.ndimension() == 3 else None
 
-        n = None
-        num_vecs = None
-        if self.ndimension() == 3:
-            _, n, num_vecs = left_vecs.size()
-        else:
-            n, num_vecs = left_vecs.size()
+        *batch_shape, n, num_vecs = left_vecs.size()
 
         if isinstance(self.right_lazy_tensor, RootLazyTensor):
             right_root = self.right_lazy_tensor.root.evaluate()
             left_factor = left_vecs.unsqueeze(-2) * right_root.unsqueeze(-1)
             right_factor = right_vecs.unsqueeze(-2) * right_root.unsqueeze(-1)
-
             right_rank = right_root.size(-1)
         else:
             right_rank = n
@@ -162,19 +158,14 @@ class MulLazyTensor(LazyTensor):
             left_factor = left_vecs.unsqueeze(-2) * self.right_lazy_tensor.evaluate().unsqueeze(-1)
             right_factor = right_vecs.unsqueeze(-2) * eye.unsqueeze(-1)
 
-        if batch_size is None:
-            left_factor = left_factor.view(n, num_vecs * right_rank)
-            right_factor = right_factor.view(n, num_vecs * right_rank)
-        else:
-            left_factor = left_factor.view(batch_size, n, num_vecs * right_rank)
-            right_factor = right_factor.view(batch_size, n, num_vecs * right_rank)
+        left_factor = left_factor.view(*batch_shape, n, num_vecs * right_rank)
+        right_factor = right_factor.view(*batch_shape, n, num_vecs * right_rank)
         left_deriv_args = self.left_lazy_tensor._quad_form_derivative(left_factor, right_factor)
 
         if isinstance(self.left_lazy_tensor, RootLazyTensor):
             left_root = self.left_lazy_tensor.root.evaluate()
             left_factor = left_vecs.unsqueeze(-2) * left_root.unsqueeze(-1)
             right_factor = right_vecs.unsqueeze(-2) * left_root.unsqueeze(-1)
-
             left_rank = left_root.size(-1)
         else:
             left_rank = n
@@ -182,12 +173,8 @@ class MulLazyTensor(LazyTensor):
             left_factor = left_vecs.unsqueeze(-2) * self.left_lazy_tensor.evaluate().unsqueeze(-1)
             right_factor = right_vecs.unsqueeze(-2) * eye.unsqueeze(-1)
 
-        if batch_size is None:
-            left_factor = left_factor.view(n, num_vecs * left_rank)
-            right_factor = right_factor.view(n, num_vecs * left_rank)
-        else:
-            left_factor = left_factor.view(batch_size, n, num_vecs * left_rank)
-            right_factor = right_factor.view(batch_size, n, num_vecs * left_rank)
+        left_factor = left_factor.view(*batch_shape, n, num_vecs * left_rank)
+        right_factor = right_factor.view(*batch_shape, n, num_vecs * left_rank)
         right_deriv_args = self.right_lazy_tensor._quad_form_derivative(left_factor, right_factor)
 
         return tuple(list(left_deriv_args) + list(right_deriv_args))
diff --git a/gpytorch/lazy/toeplitz_lazy_tensor.py b/gpytorch/lazy/toeplitz_lazy_tensor.py
index 91a9c410..2c4d451c 100644
--- a/gpytorch/lazy/toeplitz_lazy_tensor.py
+++ b/gpytorch/lazy/toeplitz_lazy_tensor.py
@@ -22,16 +22,16 @@ class ToeplitzLazyTensor(LazyTensor):
             left_vecs = left_vecs.unsqueeze(1)
             right_vecs = right_vecs.unsqueeze(1)
 
-        res = (sym_toeplitz_derivative_quadratic_form(left_vecs, right_vecs),)
-        if self.column.ndimension() == 1 and res[0].ndimension() == 2:
-            res = (res[0].sum(0),)
-        return res
+        res = sym_toeplitz_derivative_quadratic_form(left_vecs, right_vecs)
+
+        # Collapse any expanded broadcast dimensions
+        if res.dim() > self.column.dim():
+            res = res.view(-1, *self.column.shape).sum(0)
+
+        return res,
 
     def _size(self):
-        if self.column.ndimension() == 2:
-            return torch.Size((self.column.size(0), self.column.size(-1), self.column.size(-1)))
-        else:
-            return torch.Size((self.column.size(-1), self.column.size(-1)))
+        return torch.Size((*self.column.shape, self.column.size(-1)))
 
     def _transpose_nonbatch(self):
         return ToeplitzLazyTensor(self.column)
@@ -50,7 +50,7 @@ class ToeplitzLazyTensor(LazyTensor):
         """
         Gets the diagonal of the Toeplitz matrix wrapped by this object.
         """
-        diag_term = self.column.select(-1, 0)
+        diag_term = self.column[..., 0]
         if self.column.ndimension() > 1:
             diag_term = diag_term.unsqueeze(-1)
         return diag_term.expand(*self.column.size())
diff --git a/gpytorch/utils/__init__.py b/gpytorch/utils/__init__.py
index b409d94f..dde02c08 100644
--- a/gpytorch/utils/__init__.py
+++ b/gpytorch/utils/__init__.py
@@ -3,6 +3,8 @@
 from .memoize import cached
 from .linear_cg import linear_cg
 from .stochastic_lq import StochasticLQ
+from . import batch
+from . import broadcasting
 from . import cholesky
 from . import eig
 from . import fft
@@ -26,6 +28,8 @@ def prod(items):
 
 
 __all__ = [
+    "batch",
+    "broadcasting",
     "cached",
     "linear_cg",
     "StochasticLQ",
diff --git a/gpytorch/utils/batch.py b/gpytorch/utils/batch.py
new file mode 100644
index 00000000..5e65f6f9
--- /dev/null
+++ b/gpytorch/utils/batch.py
@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+
+import torch
+
+
+def _create_batch_indices(batch_shape, repeat_size, device):
+    """
+    A helper method which creates a list of batch indices for use with the LazyTensor _get_indices method
+    """
+    batch_indices = []
+    for i, batch_size in enumerate(batch_shape):
+        batch_index = torch.arange(0, batch_size, dtype=torch.long, device=device).unsqueeze(-1)
+        batch_index = batch_index.repeat(
+            torch.Size(batch_shape[:i]).numel(),
+            torch.Size(batch_shape[i + 1:]).numel() * repeat_size
+        ).view(-1)
+        batch_indices.append(batch_index)
+    return batch_indices
diff --git a/gpytorch/utils/interpolation.py b/gpytorch/utils/interpolation.py
index 98376e71..03f1d324 100644
--- a/gpytorch/utils/interpolation.py
+++ b/gpytorch/utils/interpolation.py
@@ -192,7 +192,9 @@ def left_t_interp(interp_indices, interp_values, rhs, output_dim):
     # Define a bunch of sizes
     num_data, num_interp = interp_values.shape[-2:]
     num_cols = rhs.size(-1)
-    batch_shape = torch.Size(values.shape[:-3])
+    interp_shape = torch.Size((*interp_indices.shape[:-2], output_dim, num_data))
+    output_shape = _matmul_broadcast_shape(interp_shape, rhs.shape)
+    batch_shape = output_shape[:-2]
     batch_size = batch_shape.numel()
 
     # Using interp_indices, create a sparse matrix that will sum up the values
diff --git a/gpytorch/utils/sparse.py b/gpytorch/utils/sparse.py
index 36eb041f..c61fe8b3 100644
--- a/gpytorch/utils/sparse.py
+++ b/gpytorch/utils/sparse.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 
-from operator import mul
-
 import torch
+from .broadcasting import _matmul_broadcast_shape
 
 
 def make_sparse_from_indices_and_values(interp_indices, interp_values, num_rows):
@@ -84,10 +83,34 @@ def bdsmm(sparse, dense):
     """
     Batch dense-sparse matrix multiply
     """
+    # Make the batch sparse matrix into a block-diagonal matrix
     if sparse.ndimension() > 2:
-        batch_size, num_rows, num_cols = sparse.size()
-        batch_assignment = sparse._indices()[0]
-        indices = sparse._indices()[1:].clone()
+        # Expand the tensors to account for broadcasting
+        output_shape = _matmul_broadcast_shape(sparse.shape, dense.shape)
+        expanded_sparse_shape = output_shape[:-2] + sparse.shape[-2:]
+        unsqueezed_sparse_shape = [1 for _ in range(len(output_shape) - sparse.dim())] + list(sparse.shape)
+        repeat_sizes = tuple(
+            output_size // sparse_size
+            for output_size, sparse_size in zip(expanded_sparse_shape, unsqueezed_sparse_shape)
+        )
+        sparse = sparse_repeat(sparse, *repeat_sizes)
+        dense = dense.expand(*output_shape[:-2], dense.size(-2), dense.size(-1))
+
+        # Figure out how much need to be added to the row/column indices to create
+        # a block-diagonal matrix
+        *batch_shape, num_rows, num_cols = sparse.shape
+        batch_size = torch.Size(batch_shape).numel()
+        batch_multiplication_factor = torch.tensor(
+            [torch.Size(batch_shape[i + 1:]).numel() for i in range(len(batch_shape))],
+            dtype=torch.long, device=sparse.device
+        )
+        if batch_multiplication_factor.is_cuda:
+            batch_assignment = (sparse._indices()[:-2].float().t() @ batch_multiplication_factor.float()).long()
+        else:
+            batch_assignment = sparse._indices()[:-2].t() @ batch_multiplication_factor
+
+        # Create block-diagonal sparse tensor
+        indices = sparse._indices()[-2:].clone()
         indices[0].add_(num_rows, batch_assignment)
         indices[1].add_(num_cols, batch_assignment)
         sparse_2d = torch.sparse_coo_tensor(
@@ -98,18 +121,20 @@ def bdsmm(sparse, dense):
             device=sparse._values().device,
         )
 
-        if dense.size(0) == 1:
-            dense = dense.repeat(batch_size, 1, 1)
         dense_2d = dense.contiguous().view(batch_size * num_cols, -1)
         res = torch.dsmm(sparse_2d, dense_2d)
-        res = res.view(batch_size, num_rows, -1)
+        res = res.view(*batch_shape, num_rows, -1)
         return res
-    elif dense.ndimension() == 3:
-        batch_size, _, num_cols = dense.size()
+
+    elif dense.dim() > 2:
+        *batch_shape, num_rows, num_cols = dense.size()
+        batch_size = torch.Size(batch_shape).numel()
+        dense = dense.view(batch_size, num_rows, num_cols)
         res = torch.dsmm(sparse, dense.transpose(0, 1).contiguous().view(-1, batch_size * num_cols))
         res = res.view(-1, batch_size, num_cols)
-        res = res.transpose(0, 1).contiguous()
+        res = res.transpose(0, 1).contiguous().view(*batch_shape, -1, num_cols)
         return res
+
     else:
         return torch.dsmm(sparse, dense)
 
@@ -186,40 +211,36 @@ def sparse_getitem(sparse, idxs):
 def sparse_repeat(sparse, *repeat_sizes):
     """
     """
-    orig_ndim = sparse.ndimension()
-    new_ndim = len(repeat_sizes)
-    orig_nvalues = sparse._indices().size(1)
-
-    # Expand the number of dimensions to match repeat_sizes
-    indices = torch.cat(
-        [
-            torch.zeros(new_ndim - orig_ndim, orig_nvalues, dtype=torch.long, device=sparse._indices().device),
-            sparse._indices(),
-        ]
-    )
-    values = sparse._values()
-    size = [1] * (new_ndim - orig_ndim) + list(sparse.size())
-
-    # Expand each dimension
-    new_indices = torch.zeros(
-        indices.size(0), indices.size(1) * mul(*repeat_sizes), dtype=indices.dtype, device=indices.device
-    )
-    new_values = values.repeat(mul(*repeat_sizes))
-    new_size = [dim_size * repeat_size for dim_size, repeat_size in zip(size, repeat_sizes)]
-
-    # Fill in new indices
-    new_indices[:, :orig_nvalues].copy_(indices)
-    unit_size = orig_nvalues
-    for i in range(new_ndim)[::-1]:
-        repeat_size = repeat_sizes[i]
-        for j in range(1, repeat_size):
-            new_indices[:, unit_size * j : unit_size * (j + 1)].copy_(new_indices[:, :unit_size])
-            new_indices[i, unit_size * j : unit_size * (j + 1)] += j * size[i]
-        unit_size *= repeat_size
-
-    return torch.sparse_coo_tensor(
-        new_indices, new_values, torch.Size(new_size), dtype=new_values.dtype, device=new_values.device
-    )
+    if len(repeat_sizes) == 1 and isinstance(repeat_sizes, tuple):
+        repeat_sizes = repeat_sizes[0]
+
+    if len(repeat_sizes) > len(sparse.shape):
+        num_new_dims = len(repeat_sizes) - len(sparse.shape)
+        new_indices = sparse._indices()
+        new_indices = torch.cat([
+            torch.zeros(num_new_dims, new_indices.size(1), dtype=new_indices.dtype, device=new_indices.device),
+            new_indices
+        ], 0)
+        sparse = torch.sparse_coo_tensor(
+            new_indices, sparse._values(),
+            torch.Size((*[1 for _ in range(num_new_dims)], *sparse.shape)),
+            dtype=sparse.dtype, device=sparse.device
+        )
+
+    for i, repeat_size in enumerate(repeat_sizes):
+        if repeat_size > 1:
+            new_indices = sparse._indices().repeat(1, repeat_size)
+            adding_factor = torch.arange(
+                0, repeat_size, dtype=new_indices.dtype, device=new_indices.device
+            ).unsqueeze_(1)
+            new_indices[i].view(repeat_size, -1).add_(adding_factor)
+            sparse = torch.sparse_coo_tensor(
+                new_indices, sparse._values().repeat(repeat_size),
+                torch.Size((*sparse.shape[:i], repeat_size * sparse.size(i), *sparse.shape[i + 1:])),
+                dtype=sparse.dtype, device=sparse.device
+            )
+
+    return sparse
 
 
 def to_sparse(dense):
diff --git a/gpytorch/utils/toeplitz.py b/gpytorch/utils/toeplitz.py
index c44280f9..89ec4198 100644
--- a/gpytorch/utils/toeplitz.py
+++ b/gpytorch/utils/toeplitz.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 import torch
-from ..utils import fft
+from ..utils import fft, broadcasting
 
 
 def toeplitz(toeplitz_column, toeplitz_row):
@@ -102,32 +102,15 @@ def toeplitz_matmul(toeplitz_column, toeplitz_row, tensor):
     if toeplitz_column.size() != toeplitz_row.size():
         raise RuntimeError("c and r should have the same length (Toeplitz matrices are necessarily square).")
 
-    is_batch = True
-    if toeplitz_column.ndimension() == 1:
-        toeplitz_column = toeplitz_column.unsqueeze(0)
-        toeplitz_row = toeplitz_row.unsqueeze(0)
-        if tensor.ndimension() < 3:
-            tensor = tensor.unsqueeze(0)
-        is_batch = False
+    toeplitz_shape = torch.Size((*toeplitz_column.shape, toeplitz_row.size(-1)))
+    output_shape = broadcasting._matmul_broadcast_shape(toeplitz_shape, tensor.shape)
+    broadcasted_t_shape = output_shape[:-1] if tensor.dim() > 1 else output_shape
 
-    if toeplitz_column.ndimension() != 2:
-        raise RuntimeError(
-            "The first two inputs to ToeplitzMV should be vectors \
-                            (or matrices, representing batch) \
-                            (first column c and row r of the Toeplitz matrix)"
-        )
-
-    if toeplitz_column.size(0) == 1:
-        toeplitz_column = toeplitz_column.expand(tensor.size(0), toeplitz_column.size(1))
-        toeplitz_row = toeplitz_row.expand(tensor.size(0), toeplitz_row.size(1))
-
-    if toeplitz_column.size()[:2] != tensor.size()[:2]:
-        raise RuntimeError(
-            "Dimension mismatch: attempting to multiply a {}x{} Toeplitz matrix "
-            "against a matrix with leading dimension {}.".format(
-                len(toeplitz_column), len(toeplitz_column), len(tensor)
-            )
-        )
+    if tensor.ndimension() == 1:
+        tensor = tensor.unsqueeze(-1)
+    toeplitz_column = toeplitz_column.expand(*broadcasted_t_shape).contiguous().view(-1, toeplitz_column.size(-1))
+    toeplitz_row = toeplitz_row.expand(*broadcasted_t_shape).contiguous().view(-1, toeplitz_row.size(-1))
+    tensor = tensor.expand(*output_shape).contiguous().view(-1, *tensor.shape[-2:])
 
     if not torch.equal(toeplitz_column[:, 0], toeplitz_row[:, 0]):
         raise RuntimeError(
@@ -139,44 +122,31 @@ def toeplitz_matmul(toeplitz_column, toeplitz_row, tensor):
     if type(toeplitz_column) != type(toeplitz_row) or type(toeplitz_column) != type(tensor):
         raise RuntimeError("The types of all inputs to ToeplitzMV must match.")
 
-    output_dims = tensor.ndimension()
-    if output_dims == 2:
-        tensor = tensor.unsqueeze(2)
-
-    if toeplitz_column.size(1) == 1:
-        output = toeplitz_column.view(-1, 1, 1).matmul(tensor)
-
-    else:
-        batch_size, orig_size, num_rhs = tensor.size()
-        r_reverse = toeplitz_row[:, 1:].flip(dims=(1,))
+    batch_size, orig_size, num_rhs = tensor.size()
+    r_reverse = toeplitz_row[:, 1:].flip(dims=(1,))
 
-        c_r_rev = torch.zeros(batch_size, orig_size + r_reverse.size(1), dtype=tensor.dtype, device=tensor.device)
-        c_r_rev[:, :orig_size] = toeplitz_column
-        c_r_rev[:, orig_size:] = r_reverse
+    c_r_rev = torch.zeros(batch_size, orig_size + r_reverse.size(1), dtype=tensor.dtype, device=tensor.device)
+    c_r_rev[:, :orig_size] = toeplitz_column
+    c_r_rev[:, orig_size:] = r_reverse
 
-        temp_tensor = torch.zeros(
-            batch_size, 2 * orig_size - 1, num_rhs, dtype=toeplitz_column.dtype, device=toeplitz_column.device
-        )
-        temp_tensor[:, :orig_size, :] = tensor
-
-        fft_M = fft.fft1(temp_tensor.transpose(1, 2).contiguous())
-        fft_c = fft.fft1(c_r_rev).unsqueeze(1).expand_as(fft_M)
-        fft_product = torch.zeros_like(fft_M)
-
-        fft_product[:, :, :, 0].addcmul_(fft_c[:, :, :, 0], fft_M[:, :, :, 0])
-        fft_product[:, :, :, 0].addcmul_(-1, fft_c[:, :, :, 1], fft_M[:, :, :, 1])
-        fft_product[:, :, :, 1].addcmul_(fft_c[:, :, :, 1], fft_M[:, :, :, 0])
-        fft_product[:, :, :, 1].addcmul_(fft_c[:, :, :, 0], fft_M[:, :, :, 1])
+    temp_tensor = torch.zeros(
+        batch_size, 2 * orig_size - 1, num_rhs, dtype=toeplitz_column.dtype, device=toeplitz_column.device
+    )
+    temp_tensor[:, :orig_size, :] = tensor
 
-        output = fft.ifft1(fft_product).transpose(1, 2)
-        output = output[:, :orig_size, :]
+    fft_M = fft.fft1(temp_tensor.transpose(1, 2).contiguous())
+    fft_c = fft.fft1(c_r_rev).unsqueeze(1).expand_as(fft_M)
+    fft_product = torch.zeros_like(fft_M)
 
-    if output_dims == 2:
-        output = output.squeeze(2)
+    fft_product[:, :, :, 0].addcmul_(fft_c[:, :, :, 0], fft_M[:, :, :, 0])
+    fft_product[:, :, :, 0].addcmul_(-1, fft_c[:, :, :, 1], fft_M[:, :, :, 1])
+    fft_product[:, :, :, 1].addcmul_(fft_c[:, :, :, 1], fft_M[:, :, :, 0])
+    fft_product[:, :, :, 1].addcmul_(fft_c[:, :, :, 0], fft_M[:, :, :, 1])
 
-    if not is_batch:
-        output = output.squeeze(0)
+    output = fft.ifft1(fft_product).transpose(1, 2)
+    output = output[:, :orig_size, :]
 
+    output = output.view(*output_shape)
     return output
 
 
@@ -215,34 +185,21 @@ def sym_toeplitz_derivative_quadratic_form(left_vectors, right_vectors):
         left_vectors = left_vectors.unsqueeze(1)
         right_vectors = right_vectors.unsqueeze(1)
 
-    left_vectors = left_vectors.transpose(-1, -2)
-    right_vectors = right_vectors.transpose(-1, -2)
+    batch_shape = left_vectors.shape[:-2]
+    toeplitz_size = left_vectors.size(-2)
+    num_vectors = left_vectors.size(-1)
 
-    if left_vectors.ndimension() == 3:
-        batch = True
-        batch_size, s, _ = left_vectors.size()
-        left_vectors = left_vectors.contiguous().view(batch_size * s, -1)
-        right_vectors = right_vectors.contiguous().view(batch_size * s, -1)
-    else:
-        batch = False
-
-    s, m = left_vectors.size()
+    left_vectors = left_vectors.transpose(-1, -2).contiguous()
+    right_vectors = right_vectors.transpose(-1, -2).contiguous()
 
-    left_vectors.contiguous()
-    right_vectors.contiguous()
+    columns = torch.zeros_like(left_vectors)
+    columns[..., 0] = left_vectors[..., 0]
+    res = toeplitz_matmul(columns, left_vectors, right_vectors.unsqueeze(-1))
+    rows = left_vectors.flip(dims=(-1,))
+    columns[..., 0] = rows[..., 0]
+    res += toeplitz_matmul(columns, rows, torch.flip(right_vectors, dims=(-1,)).unsqueeze(-1))
 
-    columns = torch.zeros(s, m, dtype=left_vectors.dtype, device=left_vectors.device)
-    columns[:, 0] = left_vectors[:, 0]
-    res = toeplitz_matmul(columns, left_vectors, right_vectors)
-    rows = left_vectors.flip(dims=(1,))
-    columns[:, 0] = rows[:, 0]
-    res += toeplitz_matmul(columns, rows, torch.flip(right_vectors, dims=(1,)))
-
-    if not batch:
-        res = res.sum(0)
-        res[0] -= (left_vectors * right_vectors).sum()
-    else:
-        res = res.contiguous().view(batch_size, -1, m).sum(1)
-        res[:, 0] -= (left_vectors * right_vectors).view(batch_size, -1).sum(1)
+    res = res.contiguous().view(*batch_shape, num_vectors, toeplitz_size).sum(-2)
+    res[..., 0] -= (left_vectors * right_vectors).view(*batch_shape, -1).sum(-1)
 
     return res


Test patch:
diff --git a/test/functions/test_dsmm.py b/test/functions/test_dsmm.py
index a6a97dab..8db9d945 100644
--- a/test/functions/test_dsmm.py
+++ b/test/functions/test_dsmm.py
@@ -26,6 +26,18 @@ class TestDSMM(unittest.TestCase):
         actual = torch.matmul(sparse.to_dense(), dense)
         self.assertLess(torch.norm(res - actual), 1e-5)
 
+    def test_forward_multi_batch(self):
+        i = torch.tensor(
+            [[0, 1, 1, 0, 0, 1], [0, 0, 0, 1, 1, 1], [0, 1, 1, 0, 1, 1], [2, 0, 2, 2, 0, 2]], dtype=torch.long
+        )
+        v = torch.tensor([3, 4, 5, 6, 7, 8], dtype=torch.float)
+        sparse = torch.sparse.FloatTensor(i, v, torch.Size([2, 2, 2, 3]))
+        dense = torch.randn(2, 2, 3, 3)
+
+        res = gpytorch.dsmm(sparse, dense)
+        actual = torch.matmul(sparse.to_dense(), dense)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
     def test_backward(self):
         i = torch.tensor([[0, 1, 1], [2, 0, 2]], dtype=torch.long)
         v = torch.tensor([3, 4, 5], dtype=torch.float)
@@ -54,6 +66,85 @@ class TestDSMM(unittest.TestCase):
         actual.backward(grad_output)
         self.assertLess(torch.norm(dense.grad - dense_copy.grad).item(), 1e-5)
 
+    def test_backward_multi_batch(self):
+        i = torch.tensor(
+            [[0, 1, 1, 0, 0, 1], [0, 0, 0, 1, 1, 1], [0, 1, 1, 0, 1, 1], [2, 0, 2, 2, 0, 2]], dtype=torch.long
+        )
+        v = torch.tensor([3, 4, 5, 6, 7, 8], dtype=torch.float)
+        sparse = torch.sparse.FloatTensor(i, v, torch.Size([2, 2, 2, 3]))
+        dense = torch.randn(2, 2, 3, 4, requires_grad=True)
+        dense_copy = dense.clone().detach().requires_grad_(True)
+        grad_output = torch.randn(2, 2, 2, 4)
+
+        res = gpytorch.dsmm(sparse, dense)
+        res.backward(grad_output)
+        actual = torch.matmul(sparse.to_dense(), dense_copy)
+        actual.backward(grad_output)
+        self.assertLess(torch.norm(dense.grad - dense_copy.grad).item(), 1e-5)
+
+    def test_broadcast_rhs(self):
+        i = torch.tensor([[0, 1, 1, 0, 1, 1], [2, 0, 2, 2, 0, 2]], dtype=torch.long)
+        v = torch.tensor([3, 4, 5, 6, 7, 8], dtype=torch.float)
+        sparse = torch.sparse.FloatTensor(i, v, torch.Size([2, 3]))
+        dense = torch.randn(4, 2, 3, 4, requires_grad=True)
+        dense_copy = dense.clone().detach().requires_grad_(True)
+
+        res = gpytorch.dsmm(sparse, dense)
+        actual = torch.matmul(sparse.to_dense(), dense_copy)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        grad_output = torch.randn(4, 2, 2, 4)
+        res.backward(grad_output)
+        actual.backward(grad_output)
+        self.assertLess(torch.norm(dense.grad - dense_copy.grad).item(), 1e-5)
+
+        i = torch.tensor([[0, 0, 0, 1, 1, 1], [0, 1, 1, 0, 1, 1], [2, 0, 2, 2, 0, 2]], dtype=torch.long)
+        v = torch.tensor([3, 4, 5, 6, 7, 8], dtype=torch.float)
+        sparse = torch.sparse.FloatTensor(i, v, torch.Size([2, 2, 3]))
+        dense = torch.randn(4, 2, 3, 4, requires_grad=True)
+        dense_copy = dense.clone().detach().requires_grad_(True)
+
+        res = gpytorch.dsmm(sparse, dense)
+        actual = torch.matmul(sparse.to_dense(), dense_copy)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        grad_output = torch.randn(4, 2, 2, 4)
+        res.backward(grad_output)
+        actual.backward(grad_output)
+        self.assertLess(torch.norm(dense.grad - dense_copy.grad).item(), 1e-5)
+
+    def test_broadcast_sparse(self):
+        i = torch.tensor([[0, 0, 0, 1, 1, 1], [0, 1, 1, 0, 1, 1], [2, 0, 2, 2, 0, 2]], dtype=torch.long)
+        v = torch.tensor([3, 4, 5, 6, 7, 8], dtype=torch.float)
+        sparse = torch.sparse.FloatTensor(i, v, torch.Size([2, 2, 3]))
+        dense = torch.randn(3, 4, requires_grad=True)
+        dense_copy = dense.clone().detach().requires_grad_(True)
+
+        res = gpytorch.dsmm(sparse, dense)
+        actual = torch.matmul(sparse.to_dense(), dense_copy)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        grad_output = torch.randn(2, 2, 4)
+        res.backward(grad_output)
+        actual.backward(grad_output)
+        self.assertLess(torch.norm(dense.grad - dense_copy.grad).item(), 1e-5)
+
+    def test_broadcast_singleton(self):
+        i = torch.tensor([[0, 0, 0, 1, 1, 1], [0, 1, 1, 0, 1, 1], [2, 0, 2, 2, 0, 2]], dtype=torch.long)
+        v = torch.tensor([3, 4, 5, 6, 7, 8], dtype=torch.float)
+        sparse = torch.sparse.FloatTensor(i, v, torch.Size([2, 2, 3]))
+        dense = torch.randn(1, 3, 4, requires_grad=True)
+        dense_copy = dense.clone().detach().requires_grad_(True)
+
+        res = gpytorch.dsmm(sparse, dense)
+        actual = torch.matmul(sparse.to_dense(), dense_copy)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        grad_output = torch.randn(2, 2, 4)
+        res.backward(grad_output)
+        actual.backward(grad_output)
+        self.assertLess(torch.norm(dense.grad - dense_copy.grad).item(), 1e-5)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/lazy/_lazy_tensor_test_case.py b/test/lazy/_lazy_tensor_test_case.py
index 3fd5bb40..1864ee39 100644
--- a/test/lazy/_lazy_tensor_test_case.py
+++ b/test/lazy/_lazy_tensor_test_case.py
@@ -11,6 +11,8 @@ from test._utils import approx_equal
 
 
 class RectangularLazyTensorTestCase(object):
+    no_broadcast_tests = False
+
     @abstractmethod
     def create_lazy_tensor(self):
         raise NotImplementedError()
@@ -77,6 +79,65 @@ class RectangularLazyTensorTestCase(object):
                     ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
                 )
 
+    def test_matmul_matrix_broadcast(self):
+        if self.__class__.no_broadcast_tests:
+            return
+
+        # Right hand size has one more batch dimension
+        lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+        lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+        test_vector = torch.randn(3, *lazy_tensor.batch_shape, lazy_tensor.size(-1), 5)
+        res = lazy_tensor.matmul(test_vector)
+        actual = evaluated.matmul(test_vector)
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+        grad = torch.randn_like(res)
+        res.backward(gradient=grad)
+        actual.backward(gradient=grad)
+        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+            if arg_copy.grad is not None:
+                self.assertLess(
+                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                )
+
+        if lazy_tensor.ndimension() > 2:
+            # Right hand size has one fewer batch dimension
+            lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+            lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+            evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+            test_vector = torch.randn(*lazy_tensor.batch_shape[1:], lazy_tensor.size(-1), 5)
+            res = lazy_tensor.matmul(test_vector)
+            actual = evaluated.matmul(test_vector)
+            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+            grad = torch.randn_like(res)
+            res.backward(gradient=grad)
+            actual.backward(gradient=grad)
+            for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+                if arg_copy.grad is not None:
+                    self.assertLess(
+                        ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                    )
+
+            # Right hand size has a singleton dimension
+            lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+            lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+            evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+            test_vector = torch.randn(*lazy_tensor.batch_shape[:-1], 1, lazy_tensor.size(-1), 5)
+            res = lazy_tensor.matmul(test_vector)
+            actual = evaluated.matmul(test_vector)
+            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+            grad = torch.randn_like(res)
+            res.backward(gradient=grad)
+            actual.backward(gradient=grad)
+            for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+                if arg_copy.grad is not None:
+                    self.assertLess(
+                        ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                    )
+
     def test_evaluate(self):
         lazy_tensor = self.create_lazy_tensor()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
@@ -390,6 +451,71 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
             ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
         )
 
+    def test_inv_matmul_matrix_broadcast(self):
+        if self.__class__.no_broadcast_tests:
+            return
+
+        # Right hand size has one more batch dimension
+        lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+        lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+        test_vector = torch.randn(3, *lazy_tensor.batch_shape, lazy_tensor.size(-1), 5, requires_grad=True)
+        test_vector_copy = test_vector.clone().detach().requires_grad_(True)
+        with gpytorch.settings.max_cg_iterations(100):
+            res = lazy_tensor.inv_matmul(test_vector)
+        actual = evaluated.inverse().matmul(test_vector_copy)
+        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+        grad = torch.randn_like(res)
+        res.backward(gradient=grad)
+        actual.backward(gradient=grad)
+        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+            if arg_copy.grad is not None:
+                self.assertLess(
+                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                )
+
+        if lazy_tensor.ndimension() > 2:
+            # Right hand size has one fewer batch dimension
+            lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+            lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+            evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+            test_vector = torch.randn(*lazy_tensor.batch_shape[1:], lazy_tensor.size(-1), 5)
+            test_vector_copy = test_vector.clone().detach().requires_grad_(True)
+            with gpytorch.settings.max_cg_iterations(100):
+                res = lazy_tensor.inv_matmul(test_vector)
+            actual = evaluated.inverse().matmul(test_vector_copy)
+            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+            grad = torch.randn_like(res)
+            res.backward(gradient=grad)
+            actual.backward(gradient=grad)
+            for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+                if arg_copy.grad is not None:
+                    self.assertLess(
+                        ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                    )
+
+            # Right hand size has a singleton dimension
+            lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+            lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
+            evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+            test_vector = torch.randn(*lazy_tensor.batch_shape[:-1], 1, lazy_tensor.size(-1), 5, requires_grad=True)
+            test_vector_copy = test_vector.clone().detach().requires_grad_(True)
+            with gpytorch.settings.max_cg_iterations(100):
+                res = lazy_tensor.inv_matmul(test_vector)
+            actual = evaluated.inverse().matmul(test_vector_copy)
+            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+
+            grad = torch.randn_like(res)
+            res.backward(gradient=grad)
+            actual.backward(gradient=grad)
+            for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+                if arg_copy.grad is not None:
+                    self.assertLess(
+                        ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
+                    )
+
     def test_inv_matmul_matrix_with_left(self):
         lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
         lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
diff --git a/test/lazy/test_batch_repeat_lazy_tensor.py b/test/lazy/test_batch_repeat_lazy_tensor.py
index 40809f69..9ef9cde4 100644
--- a/test/lazy/test_batch_repeat_lazy_tensor.py
+++ b/test/lazy/test_batch_repeat_lazy_tensor.py
@@ -8,6 +8,7 @@ from test.lazy._lazy_tensor_test_case import LazyTensorTestCase
 
 class TestBatchRepeatLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 0
+    no_broadcast_tests = True
     should_test_sample = True
 
     def create_lazy_tensor(self):
@@ -21,6 +22,7 @@ class TestBatchRepeatLazyTensor(LazyTensorTestCase, unittest.TestCase):
 
 class TestBatchRepeatLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
     seed = 0
+    no_broadcast_tests = True
     should_test_sample = True
 
     def create_lazy_tensor(self):
diff --git a/test/lazy/test_block_diag_lazy_tensor.py b/test/lazy/test_block_diag_lazy_tensor.py
index 2ab0a0a1..aca20169 100644
--- a/test/lazy/test_block_diag_lazy_tensor.py
+++ b/test/lazy/test_block_diag_lazy_tensor.py
@@ -9,6 +9,7 @@ from test.lazy._lazy_tensor_test_case import LazyTensorTestCase
 class TestBlockDiagLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
+    no_broadcast_tests = True
 
     def create_lazy_tensor(self):
         blocks = torch.randn(8, 4, 4)
@@ -28,6 +29,7 @@ class TestBlockDiagLazyTensor(LazyTensorTestCase, unittest.TestCase):
 class TestBlockDiagLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
+    no_broadcast_tests = True
 
     def create_lazy_tensor(self):
         blocks = torch.randn(8, 4, 4)
diff --git a/test/lazy/test_cached_cg_lazy_tensor.py b/test/lazy/test_cached_cg_lazy_tensor.py
index 714d1df8..c7b80a8f 100644
--- a/test/lazy/test_cached_cg_lazy_tensor.py
+++ b/test/lazy/test_cached_cg_lazy_tensor.py
@@ -175,6 +175,9 @@ class TestCachedCGLazyTensorNoLogdet(LazyTensorTestCase, unittest.TestCase):
         actual = lazy_tensor.inv_matmul(test_mat)
         self.assertLess(torch.norm(res - actual) / actual.norm(), 0.1)
 
+    def test_inv_matmul_matrix_broadcast(self):
+        pass
+
 
 class TestCachedCGLazyTensor(TestCachedCGLazyTensorNoLogdet):
     seed = 0
diff --git a/test/lazy/test_cat_lazy_tensor.py b/test/lazy/test_cat_lazy_tensor.py
index e079a5a6..a5cb0531 100644
--- a/test/lazy/test_cat_lazy_tensor.py
+++ b/test/lazy/test_cat_lazy_tensor.py
@@ -9,6 +9,7 @@ from test.lazy._lazy_tensor_test_case import LazyTensorTestCase
 class TestCatLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
+    no_broadcast_tests = True
 
     def create_lazy_tensor(self):
         root = torch.randn(6, 7)
@@ -31,6 +32,7 @@ class TestCatLazyTensor(LazyTensorTestCase, unittest.TestCase):
 class TestCatLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
+    no_broadcast_tests = True
 
     def create_lazy_tensor(self):
         root = torch.randn(3, 6, 7)
diff --git a/test/lazy/test_constant_mul_lazy_tensor.py b/test/lazy/test_constant_mul_lazy_tensor.py
index 4d3a0645..c0bfdd61 100644
--- a/test/lazy/test_constant_mul_lazy_tensor.py
+++ b/test/lazy/test_constant_mul_lazy_tensor.py
@@ -16,7 +16,7 @@ class TestConstantMulLazyTensor(LazyTensorTestCase, unittest.TestCase):
         return ToeplitzLazyTensor(column) * constant
 
     def evaluate_lazy_tensor(self, lazy_tensor):
-        constant = lazy_tensor.constant
+        constant = lazy_tensor.expanded_constant
         column = lazy_tensor.base_lazy_tensor.column
         return sym_toeplitz(column) * constant
 
@@ -31,7 +31,7 @@ class TestConstantMulLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
         return ToeplitzLazyTensor(column) * constant
 
     def evaluate_lazy_tensor(self, lazy_tensor):
-        constant = lazy_tensor.constant
+        constant = lazy_tensor.expanded_constant
         column = lazy_tensor.base_lazy_tensor.column
         return torch.cat([sym_toeplitz(column[0]).unsqueeze(0), sym_toeplitz(column[1]).unsqueeze(0)]) * constant.view(
             2, 1, 1
diff --git a/test/lazy/test_sum_batch_lazy_tensor.py b/test/lazy/test_sum_batch_lazy_tensor.py
index 3c475ee3..2cf29cb2 100644
--- a/test/lazy/test_sum_batch_lazy_tensor.py
+++ b/test/lazy/test_sum_batch_lazy_tensor.py
@@ -9,6 +9,7 @@ from test.lazy._lazy_tensor_test_case import LazyTensorTestCase
 class TestSumBatchLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 6
     should_test_sample = True
+    no_broadcast_tests = True
 
     def create_lazy_tensor(self):
         blocks = torch.randn(12, 4, 4)
@@ -24,6 +25,7 @@ class TestSumBatchLazyTensor(LazyTensorTestCase, unittest.TestCase):
 class TestSumBatchLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
     seed = 6
     should_test_sample = True
+    no_broadcast_tests = True
 
     def create_lazy_tensor(self):
         blocks = torch.randn(12, 4, 4)
diff --git a/test/utils/test_sparse.py b/test/utils/test_sparse.py
index c59c26f8..0d44df59 100644
--- a/test/utils/test_sparse.py
+++ b/test/utils/test_sparse.py
@@ -2,7 +2,7 @@
 
 import torch
 import unittest
-from gpytorch.utils.sparse import sparse_eye, sparse_getitem, sparse_repeat, to_sparse
+from gpytorch.utils.sparse import sparse_eye, sparse_getitem, to_sparse
 
 
 class TestSparse(unittest.TestCase):
@@ -46,29 +46,6 @@ class TestSparse(unittest.TestCase):
         res = sparse_getitem(self.sparse, (slice(2, 4), slice(1, 3)))
         self.assertTrue(torch.equal(actual, res.to_dense()))
 
-    def test_sparse_repeat_1d(self):
-        sparse_1d = sparse_getitem(self.sparse, 1)
-        actual = sparse_1d.to_dense().repeat(3, 1)
-        res = sparse_repeat(sparse_1d, 3, 1)
-        self.assertTrue(torch.equal(actual, res.to_dense()))
-
-        actual = sparse_1d.to_dense().repeat(2, 3)
-        res = sparse_repeat(sparse_1d, 2, 3)
-        self.assertTrue(torch.equal(actual, res.to_dense()))
-
-    def test_sparse_repeat_2d(self):
-        actual = self.sparse.to_dense().repeat(3, 2)
-        res = sparse_repeat(self.sparse, 3, 2)
-        self.assertTrue(torch.equal(actual, res.to_dense()))
-
-        actual = self.sparse.to_dense().repeat(1, 2)
-        res = sparse_repeat(self.sparse, 1, 2)
-        self.assertTrue(torch.equal(actual, res.to_dense()))
-
-        actual = self.sparse.to_dense().repeat(3, 1)
-        res = sparse_repeat(self.sparse, 3, 1)
-        self.assertTrue(torch.equal(actual, res.to_dense()))
-
     def test_to_sparse(self):
         actual = self.sparse
         res = to_sparse(self.sparse.to_dense())