Repository: cornellius-gp/gpytorch
Pull request #322 (2018-10-12 20:23:44)

[WIP] LazyTensors can handle multiple batches
**Note** this PR does not break any existing code, and it passes all tests.

Proposal:
- LazyTensors (and all their operations) can handle arbitrarily many batch dimensions
- LazyTensors have a `.batch_shape` property and a `.matrix_shape` property - similar to `Distributions` (automatically determined from the `LazyTensor.shape` property)
- Merge LazyTensor._get_indices and LazyTensor._batch_get_indices 

LazyTensors right now are hard-coded to either have no batch dimension (just a matrix) or 1 batch dimension. We have lots of `if... else...` statement to deal with these cases. This PR would make it possible to make LazyTensors with arbitrarily many batch dimensions.

## What's done so far
- All core `gpytorch.functions` can handle arbitrarily many batch dimensions
- The core iterative methods (`linear_cg`, `lanczos_tridiag`, and `pivoted_cholesky`) all handle many batch dimensions
- LazyTensors have the new properties (`.batch_shape` and `.matrix_shape`).

Everything so far is mostly just refactors, but comments are welcome!
(Also this PR could be merged in as-is since it passes all tests, or it could wait until everything is done.)


Patch:
diff --git a/.travis.yml b/.travis.yml
index 33152dc1..e5185315 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -2,7 +2,7 @@ language: python
 
 branches:
   only:
-    - master
+    - beta
 
 python:
   - "3.6"
diff --git a/README.md b/README.md
index 255664bd..bd95b418 100644
--- a/README.md
+++ b/README.md
@@ -13,7 +13,7 @@ GPyTorch provides (1) significant GPU acceleration (through MVM based inference)
 
 ## Examples and Tutorials
 
-See our numerous [**examples and tutorials**](http://github.com/cornellius-gp/gpytorch/blob/master/examples) on how to construct all sorts of models in GPyTorch. These example notebooks and a walk through of GPyTorch are also available at our **ReadTheDocs page [here](https://gpytorch.readthedocs.io/en/latest/index.html)**
+Right now, the package is in alpha release, and while we believe that the interface is reasonably stable, things may change. For now, see our numerous [**examples and tutorials**](http://github.com/cornellius-gp/gpytorch/blob/master/examples) on how to construct all sorts of models in GPyTorch. These example notebooks and a walk through of GPyTorch are also available at our **ReadTheDocs page [here](https://gpytorch.readthedocs.io/en/latest/index.html)**
 
 ## Installation
 
@@ -81,18 +81,20 @@ pip install flake8  # if not already installed
 flake8
 ```
 
-## The Team
+## Founding Team
 
-GPyTorch is primarily maintained by:
-- [Jake Gardner](http://github.com/jacobrgardner) (Cornell University)
-- [Geoff Pleiss](http://github.com/gpleiss) (Cornell University)
-- [Kilian Weinberger](http://kilian.cs.cornell.edu/) (Cornell University)
-- [Andrew Gordon Wilson](https://people.orie.cornell.edu/andrew/) (Cornell University)
-- [Max Balandat](https://research.fb.com/people/balandat-max/) (Facebook)
+GPyTorch is developed at Cornell University by
+- [Jake Gardner](http://github.com/jacobrgardner) (lead developer)
+- [Geoff Pleiss](http://github.com/gpleiss) (lead developer)
+- [Kilian Weinberger](http://kilian.cs.cornell.edu/)
+- [Andrew Gordon Wilson](https://people.orie.cornell.edu/andrew/)
+- Max Balandat
+- Eytan Bakshy
+- David Arbour
 
 <img width="300" src=https://brand.cornell.edu/assets/images/downloads/logos/cornell_logo_simple/cornell_logo_simple.svg alt="Cornell Logo" />
 
-We would like to thank our other contributors including (but not limited to) Eytan Bakshy, David Arbour, Ruihan Wu, Bram Wallace, Sam Stanton, and Jared Frank.
+We would like to thank our other contributors including (but not limited to) Ruihan Wu, Bram Wallace, Sam Stanton, and Jared Frank.
 
 ## Acknowledgements
 Development of GPyTorch is supported by funding from [Facebook](https://research.fb.com/) and the [Bill and Melinda Gates Foundation](https://www.gatesfoundation.org/).
diff --git a/environment.yml b/environment.yml
index 61b97370..783a0c80 100644
--- a/environment.yml
+++ b/environment.yml
@@ -2,5 +2,5 @@ name: gpytorch
 channels:
   - pytorch
 dependencies:
-  - pytorch-nightly
+  - pytorch>=0.4.1
   - torchvision
diff --git a/environment_cuda90.yml b/environment_cuda90.yml
new file mode 100644
index 00000000..df0b1a10
--- /dev/null
+++ b/environment_cuda90.yml
@@ -0,0 +1,7 @@
+name: gpytorch
+channels:
+  - pytorch
+dependencies:
+  - pytorch>=0.4.1
+  - torchvision
+  - cuda90
diff --git a/examples/05_Scalable_GP_Regression_Multidimensional/KISSGP_Deep_Kernel_Regression_With_Stochastic_Variational_Inference_CUDA.ipynb b/examples/05_Scalable_GP_Regression_Multidimensional/KISSGP_Deep_Kernel_Regression_With_Stochastic_Variational_Inference_CUDA.ipynb
index ed469ed7..7ad31413 100644
--- a/examples/05_Scalable_GP_Regression_Multidimensional/KISSGP_Deep_Kernel_Regression_With_Stochastic_Variational_Inference_CUDA.ipynb
+++ b/examples/05_Scalable_GP_Regression_Multidimensional/KISSGP_Deep_Kernel_Regression_With_Stochastic_Variational_Inference_CUDA.ipynb
@@ -61,7 +61,7 @@
     "y = data[:, -1]\n",
     "\n",
     "# Use the first 80% of the data for training, and the last 20% for testing.\n",
-    "train_n = int(floor(0.8*len(X)))\n",
+    "train_n = int(floor(0.9*len(X)))\n",
     "\n",
     "train_x = X[:train_n, :].contiguous().cuda()\n",
     "train_y = y[:train_n].contiguous().cuda()\n",
diff --git a/examples/05_Scalable_GP_Regression_Multidimensional/Scalable_Kernel_Interpolation_for_Products_CUDA.ipynb b/examples/05_Scalable_GP_Regression_Multidimensional/Scalable_Kernel_Interpolation_for_Products_CUDA.ipynb
deleted file mode 100644
index 621bbdd3..00000000
--- a/examples/05_Scalable_GP_Regression_Multidimensional/Scalable_Kernel_Interpolation_for_Products_CUDA.ipynb
+++ /dev/null
@@ -1,308 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Scalable Kernel Interpolation for Product Kernels (SKIP)\n",
-    "\n",
-    "## Overview\n",
-    "\n",
-    "In this notebook, we'll overview of how to use SKIP, a method that exploits product structure in some kernels to reduce the dependency of SKI on the data dimensionality from exponential to linear. \n",
-    "\n",
-    "The most important practical consideration to note in this notebook is the use of `gpytorch.settings.max_root_decomposition_size`, which we explain the use of right before the training loop cell."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/jrg365/anaconda3/lib/python3.6/site-packages/matplotlib/__init__.py:465: UserWarning: matplotlibrc text.usetex option can not be used unless TeX is installed on your system\n",
-      "  warnings.warn('matplotlibrc text.usetex option can not be used unless '\n",
-      "/home/jrg365/anaconda3/lib/python3.6/site-packages/matplotlib/__init__.py:473: UserWarning: matplotlibrc text.usetex can not be used with *Agg backend unless dvipng-1.6 or later is installed on your system\n",
-      "  'your system' % dvipng_req)\n"
-     ]
-    }
-   ],
-   "source": [
-    "import math\n",
-    "import torch\n",
-    "import gpytorch\n",
-    "from matplotlib import pyplot as plt\n",
-    "\n",
-    "# Make plots inline\n",
-    "%matplotlib inline"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Loading Data\n",
-    "\n",
-    "For this example notebook, we'll be using the `elevators` UCI dataset used in the paper. Running the next cell downloads a copy of the dataset that has already been scaled and normalized appropriately. For this notebook, we'll simply be splitting the data using the first 80% of the data as training and the last 20% as testing.\n",
-    "\n",
-    "**Note**: Running the next cell will attempt to download a ~400 KB dataset file to the current directory."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import urllib.request\n",
-    "import os.path\n",
-    "from scipy.io import loadmat\n",
-    "from math import floor\n",
-    "\n",
-    "if not os.path.isfile('elevators.mat'):\n",
-    "    print('Downloading \\'elevators\\' UCI dataset...')\n",
-    "    urllib.request.urlretrieve('https://drive.google.com/uc?export=download&id=1jhWL3YUHvXIaftia4qeAyDwVxo6j1alk', 'elevators.mat')\n",
-    "    \n",
-    "data = torch.Tensor(loadmat('elevators.mat')['data'])\n",
-    "X = data[:, :-1]\n",
-    "X = X - X.min(0)[0]\n",
-    "X = 2 * (X / X.max(0)[0]) - 1\n",
-    "y = data[:, -1]\n",
-    "\n",
-    "train_n = int(floor(0.8*len(X)))\n",
-    "\n",
-    "train_x = X[:train_n, :].contiguous().cuda()\n",
-    "train_y = y[:train_n].contiguous().cuda()\n",
-    "\n",
-    "test_x = X[train_n:, :].contiguous().cuda()\n",
-    "test_y = y[train_n:].contiguous().cuda()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "torch.Size([16599, 18])"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "X.size()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Defining the GP Model\n",
-    "\n",
-    "We now define the GP model. For more details on the use of GP models, see our simpler examples. This model uses a `GridInterpolationKernel` (SKI) with an RBF base kernel. To use SKIP, we make two changes:\n",
-    "\n",
-    "- First, we use only a 1 dimensional `GridInterpolationKernel` (e.g., by passing `num_dims=1`). The idea of SKIP is to use a product of 1 dimensional `GridInterpolationKernel`s instead of a single `d` dimensional one.\n",
-    "- Next, we create a `ProductStructureKernel` that wraps our 1D `GridInterpolationKernel` with `num_dims=18`. This specifies that we want to use product structure over 18 dimensions, using the 1D `GridInterpolationKernel` in each dimension.\n",
-    "\n",
-    "**Note:** If you've explored the rest of the package, you may be wondering what the differences between `AdditiveKernel`, `AdditiveStructureKernel`, `ProductKernel`, and `ProductStructureKernel` are. The `Structure` kernels (1) assume that we want to apply a single base kernel over a fully decomposed dataset (e.g., every dimension is additive or has product structure), and (2) are significantly more efficient as a result, because they can exploit batch parallel operations instead of using for loops."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {
-    "collapsed": true
-   },
-   "outputs": [],
-   "source": [
-    "from gpytorch.means import ConstantMean\n",
-    "from gpytorch.kernels import ScaleKernel, RBFKernel, ProductStructureKernel, GridInterpolationKernel\n",
-    "from gpytorch.distributions import MultivariateNormal\n",
-    "\n",
-    "class GPRegressionModel(gpytorch.models.ExactGP):\n",
-    "    def __init__(self, train_x, train_y, likelihood):\n",
-    "        super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)\n",
-    "        self.mean_module = ConstantMean()\n",
-    "        self.base_covar_module = ScaleKernel(RBFKernel())\n",
-    "        self.covar_module = ProductStructureKernel(\n",
-    "            GridInterpolationKernel(self.base_covar_module, grid_size=100, num_dims=1), num_dims=18\n",
-    "        )\n",
-    "\n",
-    "    def forward(self, x):\n",
-    "        mean_x = self.mean_module(x)\n",
-    "        covar_x = self.covar_module(x)\n",
-    "        return MultivariateNormal(mean_x, covar_x)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "likelihood = gpytorch.likelihoods.GaussianLikelihood().cuda()\n",
-    "model = GPRegressionModel(train_x, train_y, likelihood).cuda()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Training the model\n",
-    "\n",
-    "The training loop for SKIP has one main new feature we haven't seen before: we specify the `max_root_decomposition_size`. This controls how many iterations of Lanczos we want to use for SKIP, and trades off with time and--more importantly--space. Realistically, the goal should be to set this as high as possible without running out of memory.\n",
-    "\n",
-    "In some sense, this parameter is the main trade-off of SKIP. Whereas many inducing point methods care more about the number of inducing points, because SKIP approximates one dimensional kernels, it is able to do so very well with relatively few inducing points. The main source of approximation really comes from these Lanczos decompositions we perform."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {
-    "scrolled": false
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Iter 1/25 - Loss: 0.942\n",
-      "Iter 2/25 - Loss: 0.919\n",
-      "Iter 3/25 - Loss: 0.888\n",
-      "Iter 4/25 - Loss: 0.864\n",
-      "Iter 5/25 - Loss: 0.840\n",
-      "Iter 6/25 - Loss: 0.816\n",
-      "Iter 7/25 - Loss: 0.792\n",
-      "Iter 8/25 - Loss: 0.767\n",
-      "Iter 9/25 - Loss: 0.743\n",
-      "Iter 10/25 - Loss: 0.719\n",
-      "Iter 11/25 - Loss: 0.698\n",
-      "Iter 12/25 - Loss: 0.671\n",
-      "Iter 13/25 - Loss: 0.651\n",
-      "Iter 14/25 - Loss: 0.624\n",
-      "Iter 15/25 - Loss: 0.600\n",
-      "Iter 16/25 - Loss: 0.576\n",
-      "Iter 17/25 - Loss: 0.553\n",
-      "Iter 18/25 - Loss: 0.529\n",
-      "Iter 19/25 - Loss: 0.506\n",
-      "Iter 20/25 - Loss: 0.483\n",
-      "Iter 21/25 - Loss: 0.460\n",
-      "Iter 22/25 - Loss: 0.441\n",
-      "Iter 23/25 - Loss: 0.413\n",
-      "Iter 24/25 - Loss: 0.391\n",
-      "Iter 25/25 - Loss: 0.375\n",
-      "CPU times: user 1min 6s, sys: 26.8 s, total: 1min 33s\n",
-      "Wall time: 1min 48s\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Find optimal model hyperparameters\n",
-    "model.train()\n",
-    "likelihood.train()\n",
-    "\n",
-    "# Use the adam optimizer\n",
-    "optimizer = torch.optim.SGD(model.parameters(), lr=0.1)\n",
-    "\n",
-    "# \"Loss\" for GPs - the marginal log likelihood\n",
-    "mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)\n",
-    "\n",
-    "training_iterations = 25\n",
-    "def train():\n",
-    "    for i in range(training_iterations):\n",
-    "        # Zero backprop gradients\n",
-    "        optimizer.zero_grad()\n",
-    "        with gpytorch.settings.use_toeplitz(False), gpytorch.settings.max_root_decomposition_size(30):\n",
-    "            # Get output from model\n",
-    "            output = model(train_x)\n",
-    "            # Calc loss and backprop derivatives\n",
-    "            loss = -mll(output, train_y)\n",
-    "            loss.backward()\n",
-    "        print('Iter %d/%d - Loss: %.3f' % (i + 1, training_iterations, loss.item()))\n",
-    "        optimizer.step()\n",
-    "        torch.cuda.empty_cache()\n",
-    "        \n",
-    "# See dkl_mnist.ipynb for explanation of this flag\n",
-    "with gpytorch.settings.use_toeplitz(True):\n",
-    "    %time train()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Making Predictions\n",
-    "\n",
-    "The next cell makes predictions with SKIP. We use the same max_root_decomposition size, and we also demonstrate increasing the max preconditioner size. Increasing the preconditioner size on this dataset is **not** necessary, but can make a big difference in final test performance, and is often preferable to increasing the number of CG iterations if you can afford the space."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.eval()\n",
-    "likelihood.eval()\n",
-    "with gpytorch.settings.max_preconditioner_size(10), torch.no_grad():\n",
-    "    with gpytorch.settings.use_toeplitz(False), gpytorch.settings.max_root_decomposition_size(30), gpytorch.fast_pred_var():\n",
-    "        preds = model(test_x)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Test MAE: 0.07745790481567383\n"
-     ]
-    }
-   ],
-   "source": [
-    "print('Test MAE: {}'.format(torch.mean(torch.abs(preds.mean - test_y))))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "collapsed": true
-   },
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "anaconda-cloud": {},
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.6.4"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/examples/08_Deep_Kernel_Learning/Deep_Kernel_Learning_DenseNet_CIFAR_Tutorial.ipynb b/examples/08_Deep_Kernel_Learning/Deep_Kernel_Learning_DenseNet_CIFAR_Tutorial.ipynb
index 16f374f9..8b215f33 100644
--- a/examples/08_Deep_Kernel_Learning/Deep_Kernel_Learning_DenseNet_CIFAR_Tutorial.ipynb
+++ b/examples/08_Deep_Kernel_Learning/Deep_Kernel_Learning_DenseNet_CIFAR_Tutorial.ipynb
@@ -329,7 +329,7 @@
    "source": [
     "for epoch in range(1, n_epochs + 1):\n",
     "    scheduler.step()\n",
-    "    with gpytorch.settings.use_toeplitz(False), gpytorch.settings.max_preconditioner_size(0):\n",
+    "    with gpytorch.settings.use_toeplitz(False):\n",
     "        train(epoch)\n",
     "        test()\n",
     "    state_dict = model.state_dict()\n",
diff --git a/gpytorch/functions/_inv_quad_log_det.py b/gpytorch/functions/_inv_quad_log_det.py
index 2de1501b..c49196a7 100644
--- a/gpytorch/functions/_inv_quad_log_det.py
+++ b/gpytorch/functions/_inv_quad_log_det.py
@@ -22,27 +22,23 @@ class InvQuadLogDet(Function):
     def __init__(
         self,
         representation_tree,
-        matrix_size=0,
-        batch_size=None,
-        dtype=None,
-        device=None,
+        dtype,
+        device,
+        matrix_shape,
+        batch_shape=torch.Size(),
         inv_quad=False,
         log_det=False,
         preconditioner=None,
         log_det_correction=None,
     ):
-        if not matrix_size:
-            raise RuntimeError("Matrix size must be set")
-        if device is None:
-            raise RuntimeError("device must be set")
         if not (inv_quad or log_det):
             raise RuntimeError("Either inv_quad or log_det must be true (or both)")
 
         self.representation_tree = representation_tree
-        self.matrix_size = matrix_size
-        self.batch_size = batch_size
         self.dtype = dtype
         self.device = device
+        self.matrix_shape = matrix_shape
+        self.batch_shape = batch_shape
         self.inv_quad = inv_quad
         self.log_det = log_det
         self.preconditioner = preconditioner
@@ -81,12 +77,12 @@ class InvQuadLogDet(Function):
         probe_vector_norms = None
         if self.log_det:
             num_random_probes = settings.num_trace_samples.value()
-            probe_vectors = torch.empty(self.matrix_size, num_random_probes, dtype=self.dtype, device=self.device)
+            probe_vectors = torch.empty(self.matrix_shape[-1], num_random_probes, dtype=self.dtype, device=self.device)
             probe_vectors.bernoulli_().mul_(2).add_(-1)
             probe_vector_norms = torch.norm(probe_vectors, 2, dim=-2, keepdim=True)
-            if self.batch_size is not None:
-                probe_vectors = probe_vectors.unsqueeze(0).expand(self.batch_size, self.matrix_size, num_random_probes)
-                probe_vector_norms = probe_vector_norms.unsqueeze(0).expand(self.batch_size, 1, num_random_probes)
+            if self.batch_shape is not None:
+                probe_vectors = probe_vectors.expand(*self.batch_shape, self.matrix_shape[-1], num_random_probes)
+                probe_vector_norms = probe_vector_norms.expand(*self.batch_shape, 1, num_random_probes)
             probe_vectors = probe_vectors.div(probe_vector_norms)
             rhs_list.append(probe_vectors)
 
@@ -127,12 +123,11 @@ class InvQuadLogDet(Function):
 
         # Compute log_det from tridiagonalization
         if self.log_det:
-            if self.batch_size is None:
+            if self.batch_shape is None:
                 t_mat = t_mat.unsqueeze(1)
             eigenvalues, eigenvectors = lanczos_tridiag_to_diag(t_mat)
             slq = StochasticLQ()
-            matrix_size = rhs.size(-2)
-            log_det_term, = slq.evaluate(t_mat, matrix_size, eigenvalues, eigenvectors, [lambda x: x.log()])
+            log_det_term, = slq.evaluate(self.matrix_shape, eigenvalues, eigenvectors, [lambda x: x.log()])
 
             # Add correction
             if self.log_det_correction is not None:
@@ -141,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_size is None))
+            inv_quad_term = (inv_quad_solves * inv_quad_rhs).sum(-1).sum(-1, keepdim=(self.batch_shape is None))
 
         self.num_random_probes = num_random_probes
         self.num_inv_quad_solves = num_inv_quad_solves
@@ -176,11 +171,11 @@ class InvQuadLogDet(Function):
         # Fix grad_output sizes
         if self.inv_quad:
             inv_quad_grad_output = inv_quad_grad_output.unsqueeze(-1)
-            if self.batch_size is not None:
+            if self.batch_shape is not None:
                 inv_quad_grad_output.unsqueeze_(-1)
         if compute_log_det_grad:
             log_det_grad_output = log_det_grad_output.unsqueeze(-1)
-            if self.batch_size is not None:
+            if self.batch_shape is not None:
                 log_det_grad_output.unsqueeze_(-1)
 
         # Divide up the solves
@@ -226,5 +221,4 @@ class InvQuadLogDet(Function):
             res = tuple([inv_quad_rhs_grad] + list(matrix_arg_grads))
         else:
             res = matrix_arg_grads
-
-        return tuple(res)
+        return res
diff --git a/gpytorch/functions/_matmul.py b/gpytorch/functions/_matmul.py
index 382fc972..be828c7c 100644
--- a/gpytorch/functions/_matmul.py
+++ b/gpytorch/functions/_matmul.py
@@ -12,14 +12,25 @@ class Matmul(Function):
         self.representation_tree = representation_tree
 
     def forward(self, rhs, *matrix_args):
+        orig_rhs = rhs
+
+        if rhs.ndimension() == 1:
+            is_vector = True
+            rhs = rhs.unsqueeze(-1)
+        else:
+            is_vector = False
+
         lazy_tsr = self.representation_tree(*matrix_args)
         res = lazy_tsr._matmul(rhs)
 
-        to_save = [rhs] + list(matrix_args)
+        to_save = [orig_rhs] + list(matrix_args)
         self.save_for_backward(*to_save)
         if not settings.memory_efficient.on():
             self._lazy_tsr = lazy_tsr
 
+        # Squeeze if necessary
+        if is_vector:
+            res = res.squeeze(-1)
         return res
 
     def backward(self, grad_output):
diff --git a/gpytorch/functions/_root_decomposition.py b/gpytorch/functions/_root_decomposition.py
index 10124562..a4b52cbd 100644
--- a/gpytorch/functions/_root_decomposition.py
+++ b/gpytorch/functions/_root_decomposition.py
@@ -13,24 +13,24 @@ class RootDecomposition(Function):
     def __init__(
         self,
         representation_tree,
+        max_iter,
         dtype,
         device,
-        size,
-        max_iter,
-        batch_size=None,
+        matrix_shape,
+        batch_shape,
         root=True,
         inverse=False,
-        initial_vector=None,
+        initial_vectors=None,
     ):
         self.representation_tree = representation_tree
         self.device = device
         self.dtype = dtype
-        self.size = size
+        self.matrix_shape = matrix_shape
         self.max_iter = max_iter
-        self.batch_size = batch_size
+        self.batch_shape = batch_shape
         self.root = root
         self.inverse = inverse
-        self.initial_vector = initial_vector
+        self.initial_vectors = initial_vectors
 
     def forward(self, *matrix_args):
         """
@@ -49,11 +49,11 @@ class RootDecomposition(Function):
             self.max_iter,
             dtype=self.dtype,
             device=self.device,
-            batch_size=self.batch_size,
-            n_dims=self.size,
-            init_vecs=self.initial_vector,
+            matrix_shape=self.matrix_shape,
+            batch_shape=self.batch_shape,
+            init_vecs=self.initial_vectors,
         )
-        if self.batch_size is None:
+        if self.batch_shape is None:
             q_mat = q_mat.unsqueeze(-3)
             t_mat = t_mat.unsqueeze(-3)
         if t_mat.ndimension() == 3:  # If we only used one probe vector
@@ -78,7 +78,7 @@ class RootDecomposition(Function):
         if not settings.memory_efficient.on():
             self._lazy_tsr = lazy_tsr
 
-        if self.batch_size is None:
+        if self.batch_shape is None:
             root = root.squeeze(1) if root.numel() else root
             q_mat = q_mat.squeeze(1)
             t_mat = t_mat.squeeze(1)
diff --git a/gpytorch/functions/matmul.py b/gpytorch/functions/matmul.py
deleted file mode 100644
index b48338d6..00000000
--- a/gpytorch/functions/matmul.py
+++ /dev/null
@@ -1,48 +0,0 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from __future__ import unicode_literals
-
-from torch.autograd import Function
-from .. import settings
-
-
-class Matmul(Function):
-    def __init__(self, representation_tree):
-        self.representation_tree = representation_tree
-
-    def forward(self, rhs, *matrix_args):
-        lazy_tsr = self.representation_tree(*matrix_args)
-        res = lazy_tsr._matmul(rhs)
-
-        to_save = [rhs] + list(matrix_args)
-        self.save_for_backward(*to_save)
-        if not settings.memory_efficient.on():
-            self._lazy_tsr = lazy_tsr
-
-        return res
-
-    def backward(self, grad_output):
-        rhs = self.saved_tensors[0]
-        matrix_args = self.saved_tensors[1:]
-        rhs_shape = rhs.shape
-
-        rhs_grad = None
-        arg_grads = [None] * len(matrix_args)
-
-        # input_1 gradient
-        if any(self.needs_input_grad[1:]):
-            rhs = rhs.unsqueeze(-1) if (rhs.ndimension() == 1) else rhs
-            grad_output_matrix = grad_output.unsqueeze(-1) if (grad_output.ndimension() == 1) else grad_output
-            arg_grads = self.representation_tree(*matrix_args)._quad_form_derivative(grad_output_matrix, rhs)
-
-        # input_2 gradient
-        if self.needs_input_grad[0]:
-            if hasattr(self, "_lazy_tsr"):
-                lazy_tsr = self._lazy_tsr
-            else:
-                lazy_tsr = self.representation_tree(*matrix_args)
-            rhs_grad = lazy_tsr._t_matmul(grad_output)
-            rhs_grad = rhs_grad.view(rhs_shape)
-
-        return tuple([rhs_grad] + list(arg_grads))
diff --git a/gpytorch/kernels/spectral_mixture_kernel.py b/gpytorch/kernels/spectral_mixture_kernel.py
index 0de0651a..5643a2b0 100644
--- a/gpytorch/kernels/spectral_mixture_kernel.py
+++ b/gpytorch/kernels/spectral_mixture_kernel.py
@@ -51,15 +51,19 @@ class SpectralMixtureKernel(Kernel):
             The mixture weight parameters (`b x k`).
 
     Example:
-        >>> # Non-batch
         >>> x = torch.randn(10, 5)
-        >>> covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4, ard_dum_dims=5)
+        >>> # Non-batch: Simple option
+        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
+        >>> # Non-batch: ARD (different lengthscale for each input dimension)
+        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(ard_num_dims=5))
         >>> covar = covar_module(x)  # Output: LazyVariable of size (10 x 10)
         >>>
-        >>> # Batch
         >>> batch_x = torch.randn(2, 10, 5)
-        >>> covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4, batch_size=2, ard_dum_dims=5)
-        >>> covar = covar_module(x)  # Output: LazyVariable of size (10 x 10)
+        >>> # Batch: Simple option
+        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
+        >>> # Batch: different lengthscale for each batch
+        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(batch_size=2))
+        >>> covar = covar_module(x)  # Output: LazyVariable of size (2 x 10 x 10)
 
 
     .. _Gaussian Process Kernels for Pattern Discovery and Extrapolation:
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 473e779e..51116e1e 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -31,6 +31,8 @@ class LazyTensor(object):
     In order to define a new LazyTensor class that can be used as a covariance matrix in GPyTorch, a user must define
     at a minimum the following methods (in each example, :math:`K` denotes the matrix that the LazyTensor represents)
 
+    * :func:`~gpytorch.lazy.LazyTensor._get_indices`, which returns a Tensor where the entries are determined by
+      LongTensors of indices.
     * :func:`~gpytorch.lazy.LazyTensor._matmul`, which performs a matrix multiplication :math:`KM`
     * :func:`~gpytorch.lazy.LazyTensor._quad_form_derivative`, which computes a quadratic form with the derivative,
       :math:`\mathbf{v}^{\\top}\\frac{dK}{dR}\mathbf{v}`, where :math:`R` denotes the actual tensors used to represent
@@ -39,6 +41,7 @@ class LazyTensor(object):
       :math:`\mathbf{c}`, this would return :math:`\mathbf{v}^{\\top}\\frac{dK}{d\mathbf{c}}\mathbf{v}`.
     * :func:`~gpytorch.lazy.LazyTensor._size`, which returns a :class:`torch.Size` containing the dimensions of
       :math:`K`.
+    * :func:`~gpytorch.lazy.LazyTensor._transpose_nonbatch`, which returns a transposed version of the LazyTensor
 
     In addition to these, a LazyTensor may need to define the :func:`~gpytorch.lazy.LazyTensor._transpose_nonbatch`,
     :func:`~gpytorch.lazy.LazyTensor._get_indices`, and :func:`~gpytorch.lazy.LazyTensor._batch_get_indices`
@@ -63,6 +66,75 @@ class LazyTensor(object):
         self._args = args
         self._kwargs = kwargs
 
+    def _get_indices(self, *batch_indices, left_indices, right_indices):
+        """
+        Returns entries of the matrix, indexed by batch, row, and column indices
+        """
+        raise NotImplementedError("The class {} requires a _get_indices function!".format(self.__class__.__name__))
+
+    def _matmul(self, rhs):
+        """
+        Performs a matrix multiplication :math:`KM` with the matrix :math:`K` that this LazyTensor represents. Should
+        behave as :func:`torch.matmul`. If the LazyTensor represents a batch of matrices, this method should therefore
+        operate in batch mode as well.
+
+        ..note::
+            This method is intended to be used only internally by various Functions that support backpropagation
+            (e.g., :class:`gpytorch.functions.Matmul`). Once this method is defined, it is strongly recommended that
+            one use :func:`~gpytorch.lazy.LazyTensor.matmul` instead, which makes use of this method properly.
+
+        Args:
+            rhs (:obj:`torch.tensor`): the matrix :math:`M` to multiply with.
+
+        Returns:
+            :obj:`torch.tensor`: matrix * rhs
+        """
+        raise NotImplementedError("The class {} requires a _matmul function!".format(self.__class__.__name__))
+
+    def _quad_form_derivative(self, left_vecs, right_vecs):
+        """
+        Given u (left_vecs) and v (right_vecs),
+        Computes the derivatives of (u^t K v) w.r.t. K
+
+        ..note::
+            This method is intended to be used only internally by various Functions that support backpropagation.
+            For example, this method is used internally by :func:`~gpytorch.lazy.LazyTensor.inv_quad_log_det`. It is
+            not likely that users will need to call this method directly.
+
+        Returns:
+            :obj:`torch.tensor`: derivative with respect to the arguments that are actually used to represent this
+                                   this LazyTensor.
+        """
+        raise NotImplementedError(
+            "The class {} requires a _quad_form_derivative function!".format(self.__class__.__name__)
+        )
+
+    def _size(self):
+        """
+        Returns the size of the resulting Tensor that the lazy tensor represents.
+
+        ..note::
+            This method is used internally by the related function :func:`~gpytorch.lazy.LazyTensor.size`,
+            which does some additional work. Calling this method directly is discouraged.
+
+        Returns:
+            :obj:`torch.Size`: The size of the matrix :math:`K` represented by this LazyTensor
+        """
+        raise NotImplementedError("The class {} requires a _size function!".format(self.__class__.__name__))
+
+    def _transpose_nonbatch(self):
+        """
+        Transposes non-batch dimensions (e.g. last two)
+        Implement this method, rather than transpose() or t().
+
+        ..note::
+            This method is used internally by the related function :func:`~gpytorch.lazy.LazyTensor.transpose`, which
+            does some additional work. Calling this method directly is discouraged.
+        """
+        raise NotImplementedError(
+            "The class {} requires a _transpose_nonbatch function!".format(self.__class__.__name__)
+        )
+
     @property
     def _args(self):
         return self._args_memo
@@ -71,16 +143,6 @@ class LazyTensor(object):
     def _args(self, args):
         self._args_memo = args
 
-    def _preconditioner(self):
-        """
-        (Optional) define a preconditioner (P) for linear conjugate gradients
-
-        Returns:
-            function: a function on x which performs P^{-1}(x)
-            scalar: the log determinant of P
-        """
-        return None, None
-
     def _approx_diag(self):
         """
         (Optional) returns an (approximate) diagonal of the matrix
@@ -96,6 +158,36 @@ class LazyTensor(object):
         """
         return self.diag()
 
+    def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root, test_train_covar):
+        """
+        Computes a cache for K_X*X (K_XX + sigma^2 I)^-1 K_X*X if possible. By default, this does no work and returns
+        the first argument.
+
+        Args:
+            train_train_covar_inv_root (:obj:`torch.tensor`): a root of (K_XX + sigma^2 I)^-1
+            test_train_covar (:obj:`torch.tensor`): the observed noise (from the likelihood)
+
+        Returns
+            - A precomputed cache
+        """
+        return train_train_covar_inv_root.detach()
+
+    def _exact_predictive_covar_inv_quad_form_root(self, precomputed_cache, test_train_covar):
+        """
+        Computes :math:`K_{X^{*}X} S` given a precomputed cache
+        Where :math:`S` is a tensor such that :math:`SS^{\\top} = (K_{XX} + \sigma^2 I)^{-1}`
+
+        Args:
+            precomputed_cache (:obj:`torch.tensor`): What was computed in _exact_predictive_covar_inv_quad_form_cache
+            test_train_covar (:obj:`torch.tensor`): The observed noise (from the likelihood)
+
+        Returns
+            :obj:`~gpytorch.lazy.LazyTensor`: :math:`K_{X^{*}X} S`
+        """
+        # Here the precomputed cache represents S,
+        # where S S^T = (K_XX + sigma^2 I)^-1
+        return test_train_covar.matmul(precomputed_cache)
+
     def _getitem_nonbatch(self, row_index, col_index, first_tensor_index_dim=None):
         """
         Given an index over rows and columns, gets those items from the LazyTensor.
@@ -145,24 +237,15 @@ class LazyTensor(object):
         )
         return res
 
-    def _matmul(self, rhs):
+    def _preconditioner(self):
         """
-        Performs a matrix multiplication :math:`KM` with the matrix :math:`K` that this LazyTensor represents. Should
-        behave as :func:`torch.matmul`. If the LazyTensor represents a batch of matrices, this method should therefore
-        operate in batch mode as well.
-
-        ..note::
-            This method is intended to be used only internally by various Functions that support backpropagation
-            (e.g., :class:`gpytorch.functions.Matmul`). Once this method is defined, it is strongly recommended that
-            one use :func:`~gpytorch.lazy.LazyTensor.matmul` instead, which makes use of this method properly.
-
-        Args:
-            rhs (:obj:`torch.tensor`): the matrix :math:`M` to multiply with.
+        (Optional) define a preconditioner (P) for linear conjugate gradients
 
         Returns:
-            :obj:`torch.tensor`: matrix * rhs
+            function: a function on x which performs P^{-1}(x)
+            scalar: the log determinant of P
         """
-        raise NotImplementedError("The class %s requires a _matmul function!" % self.__class__.__name__)
+        return None, None
 
     def _t_matmul(self, rhs):
         """
@@ -177,66 +260,6 @@ class LazyTensor(object):
         """
         return self.transpose(-1, -2)._matmul(rhs)
 
-    def _quad_form_derivative(self, left_vecs, right_vecs):
-        """
-        Given u (left_vecs) and v (right_vecs),
-        Computes the derivatives of (u^t K v) w.r.t. K
-
-        ..note::
-            This method is intended to be used only internally by various Functions that support backpropagation.
-            For example, this method is used internally by :func:`~gpytorch.lazy.LazyTensor.inv_quad_log_det`. It is
-            not likely that users will need to call this method directly.
-
-        Returns:
-            :obj:`torch.tensor`: derivative with respect to the arguments that are actually used to represent this
-                                   this LazyTensor.
-        """
-        raise NotImplementedError("The class %s requires a _quad_form_derivative function!" % self.__class__.__name__)
-
-    def _size(self):
-        """
-        Returns the size of the resulting Tensor that the lazy tensor represents.
-
-        ..note::
-            This method is used internally by the related function :func:`~gpytorch.lazy.LazyTensor.size`, which does
-            some additional work. Calling this method directly is discouraged.
-
-        Returns:
-            :obj:`torch.Size`: The size of the matrix :math:`K` represented by this LazyTensor
-        """
-        raise NotImplementedError("The class %s requires a _size function!" % self.__class__.__name__)
-
-    def _transpose_nonbatch(self):
-        """
-        Transposes non-batch dimensions (e.g. last two)
-
-        Implement this method, rather than transpose() or t().
-        This is because size does some additional work
-        """
-        raise NotImplementedError("The class %s requires a _transpose_nonbatch function!" % self.__class__.__name__)
-
-    def _batch_get_indices(self, batch_indices, left_indices, right_indices):
-        """
-        This function is necessary if you're using a LazyTensor as part of
-        and InterpolatedLazyTensor
-
-        Batch version of _get_indices
-        For each matrix in batch_indices, returns entries of the matrix,
-        indexed by left and right indices
-        Only works for batch lazy tensors
-        """
-        raise NotImplementedError("The class %s requires a _batch_get_indices function!" % self.__class__.__name__)
-
-    def _get_indices(self, left_indices, right_indices):
-        """
-        This function is necessary if you're using a LazyTensor as part of
-        and InterpolatedLazyTensor
-
-        Returns entries of the matrix, indexed by left and right indices
-        Only works for non-batch lazy tensors
-        """
-        raise NotImplementedError("The class %s requires a _get_indices function!" % self.__class__.__name__)
-
     def add_diag(self, diag):
         """
         Adds an element to the diagonal of the matrix.
@@ -287,6 +310,13 @@ class LazyTensor(object):
         diag = torch.tensor(jitter_val, dtype=self.dtype, device=self.device)
         return self.add_diag(diag)
 
+    @property
+    def batch_shape(self):
+        """
+        Returns the shape over which parameters are batched.
+        """
+        return torch.Size(self.shape[:-2])
+
     def clone(self):
         """
         Clones the LazyTensor (creates clones of all underlying tensors)
@@ -339,29 +369,6 @@ class LazyTensor(object):
                 new_kwargs[name] = val
         return self.__class__(*new_args, **new_kwargs)
 
-    def to(self, device_id):
-        """
-        A device-agnostic method of moving the lazy_tensor to the specified device.
-
-        Args:
-            device_id (:obj: `torch.device`): Which device to use (GPU or CPU).
-        Returns:
-            :obj:`~gpytorch.lazy.LazyTensor`: New LazyTensor identical to self on specified device
-        """
-        new_args = []
-        new_kwargs = {}
-        for arg in self._args:
-            if hasattr(arg, "to"):
-                new_args.append(arg.to(device_id))
-            else:
-                new_args.append(arg)
-        for name, val in self._kwargs.items():
-            if hasattr(val, "to"):
-                new_kwargs[name] = val.to(device_id)
-            else:
-                new_kwargs[name] = val
-        return self.__class__(*new_args, **new_kwargs)
-
     @property
     def device(self):
         return self._args[0].device
@@ -403,25 +410,15 @@ class LazyTensor(object):
         Explicitly evaluates the matrix this LazyTensor represents. This function
         should return a Tensor storing an exact representation of this LazyTensor.
         """
-        size = self.size()
-        if len(size) == 2:
-            batch_mode = False
-            num_rows, num_cols = size
-        else:
-            batch_mode = True
-            batch_size, num_rows, num_cols = size
+        num_rows, num_cols = self.matrix_shape
 
         if num_rows < num_cols:
             eye = torch.eye(num_rows, dtype=self.dtype, device=self.device)
-            if batch_mode:
-                eye = eye.unsqueeze(0).expand(batch_size, num_rows, num_rows)
-                return self.transpose(1, 2).matmul(eye).transpose(1, 2).contiguous()
-            else:
-                return self.t().matmul(eye).t().contiguous()
+            eye = eye.expand(*self.batch_shape, num_rows, num_rows)
+            return self.transpose(-1, -2).matmul(eye).transpose(-1, -2).contiguous()
         else:
             eye = torch.eye(num_cols, dtype=self.dtype, device=self.device)
-            if batch_mode:
-                eye = eye.unsqueeze(0).expand(batch_size, num_cols, num_cols)
+            eye = eye.expand(*self.batch_shape, num_cols, num_cols)
             return self.matmul(eye)
 
     def evaluate_kernel(self):
@@ -475,36 +472,6 @@ class LazyTensor(object):
         res = res + test_mean
         return res, precomputed_cache.detach()
 
-    def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root, test_train_covar):
-        """
-        Computes a cache for K_X*X (K_XX + sigma^2 I)^-1 K_X*X if possible. By default, this does no work and returns
-        the first argument.
-
-        Args:
-            train_train_covar_inv_root (:obj:`torch.tensor`): a root of (K_XX + sigma^2 I)^-1
-            test_train_covar (:obj:`torch.tensor`): the observed noise (from the likelihood)
-
-        Returns
-            - A precomputed cache
-        """
-        return train_train_covar_inv_root.detach()
-
-    def _exact_predictive_covar_inv_quad_form_root(self, precomputed_cache, test_train_covar):
-        """
-        Computes :math:`K_{X^{*}X} S` given a precomputed cache
-        Where :math:`S` is a tensor such that :math:`SS^{\\top} = (K_{XX} + \sigma^2 I)^{-1}`
-
-        Args:
-            precomputed_cache (:obj:`torch.tensor`): What was computed in _exact_predictive_covar_inv_quad_form_cache
-            test_train_covar (:obj:`torch.tensor`): The observed noise (from the likelihood)
-
-        Returns
-            :obj:`~gpytorch.lazy.LazyTensor`: :math:`K_{X^{*}X} S`
-        """
-        # Here the precomputed cache represents S,
-        # where S S^T = (K_XX + sigma^2 I)^-1
-        return test_train_covar.matmul(precomputed_cache)
-
     def exact_predictive_covar(self, num_train, likelihood, precomputed_cache=None):
         """
         Computes the posterior predictive covariance of a GP
@@ -563,15 +530,30 @@ class LazyTensor(object):
         Returns:
             - :obj:`torch.tensor` - :math:`K^{-1}M`
         """
-        # Work out batch dimension, if necessary
-        lazy_tsr = self
-        if lazy_tsr.ndimension() == 3 and tensor.ndimension() == 3:
-            if lazy_tsr.size(0) == 1 and tensor.size(0) > 1:
-                lazy_tsr = lazy_tsr.repeat(tensor.size(0), 1, 1)
-            elif tensor.size(0) == 1:
-                tensor = tensor.expand(lazy_tsr.size(0), tensor.size(1), tensor.size(2))
-        elif self.ndimension() > 3 or tensor.ndimension() > 3:
-            raise RuntimeError
+        if not self.is_square:
+            raise RuntimeError(
+                "inv_matmul only operates on (batches of) square (positive semi-definite) LazyTensors. "
+                "Got a {} of size {}.".format(self.__class__.__name__, self.size())
+            )
+
+        if (self.dim() == 2 and tensor.dim() == 1):
+            if self.shape[-1] != tensor.numel():
+                raise RuntimeError(
+                    "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                        self.shape, tensor.shape
+                    )
+                )
+        elif self.dim() != tensor.dim():
+            raise RuntimeError(
+                "LazyTensor (size={}) and right-hand-side Tensor (size={}) should have the same number "
+                "of dimensions.".format(self.shape, tensor.shape)
+            )
+        elif self.batch_shape != tensor.shape[:-2] or self.shape[-1] != tensor.shape[-2]:
+            raise RuntimeError(
+                "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                    self.shape, tensor.shape
+                )
+            )
 
         func = InvMatmul(self.representation_tree(), preconditioner=self._preconditioner()[0])
         return func(tensor, *self.representation())
@@ -606,28 +588,40 @@ class LazyTensor(object):
             - scalar - tr( tensor^T (self)^{-1} tensor )
             - scalar - log determinant
         """
-        # Work out batch dimension, if necessary
-        lazy_tsr = self
+        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 lazy_tsr.ndimension() == 3 and inv_quad_rhs.ndimension() == 3:
-                if lazy_tsr.size(0) == 1 and inv_quad_rhs.size(0) > 1:
-                    lazy_tsr = lazy_tsr.repeat(inv_quad_rhs.size(0), 1, 1)
-                elif inv_quad_rhs.size(0) == 1:
-                    inv_quad_rhs = inv_quad_rhs.expand(lazy_tsr.size(0), inv_quad_rhs.size(1), inv_quad_rhs.size(2))
-            elif self.ndimension() > 3 or inv_quad_rhs.ndimension() > 3:
-                raise RuntimeError
-
-        matrix_size = self.size(-1)
-        batch_size = self.size(0) if self.ndimension() == 3 else None
-
-        args = lazy_tsr.representation()
+            if (self.dim() == 2 and inv_quad_rhs.dim() == 1):
+                if self.shape[-1] != inv_quad_rhs.numel():
+                    raise RuntimeError(
+                        "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                            self.shape, inv_quad_rhs.shape
+                        )
+                    )
+            elif 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
+                    )
+                )
+
+        args = self.representation()
         if inv_quad_rhs is not None:
             args = [inv_quad_rhs] + list(args)
 
         res = InvQuadLogDet(
             representation_tree=self.representation_tree(),
-            matrix_size=matrix_size,
-            batch_size=batch_size,
+            matrix_shape=self.matrix_shape,
+            batch_shape=self.batch_shape,
             dtype=self.dtype,
             device=self.device,
             inv_quad=(inv_quad_rhs is not None),
@@ -637,6 +631,10 @@ class LazyTensor(object):
         )(*args)
         return res
 
+    @property
+    def is_square(self):
+        return self.matrix_shape[0] == self.matrix_shape[1]
+
     def log_det(self):
         """
         Computes an (approximate) log determinant of the matrix
@@ -663,20 +661,35 @@ class LazyTensor(object):
             is the matrix that this :obj:`gpytorch.lazy.LazyTensor` represents, and :math:`M` is the matrix input
             to this method.
         """
-
-        # Work out batch dimension, if necessary
-        lazy_tsr = self
-        if lazy_tsr.ndimension() == 3 and tensor.ndimension() == 3:
-            if lazy_tsr.size(0) == 1 and tensor.size(0) > 1:
-                lazy_tsr = lazy_tsr.repeat(tensor.size(0), 1, 1)
-            elif tensor.size(0) == 1:
-                tensor = tensor.expand(lazy_tsr.size(0), tensor.size(1), tensor.size(2))
-        elif self.ndimension() > 3 or tensor.ndimension() > 3:
-            raise RuntimeError
+        if (self.dim() == 2 and tensor.dim() == 1):
+            if self.shape[-1] != tensor.numel():
+                raise RuntimeError(
+                    "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                        self.shape, tensor.shape
+                    )
+                )
+        elif self.dim() != tensor.dim():
+            raise RuntimeError(
+                "LazyTensor (size={}) and right-hand-side Tensor (size={}) should have the same number "
+                "of dimensions.".format(self.shape, tensor.shape)
+            )
+        elif self.batch_shape != tensor.shape[:-2] or self.shape[-1] != tensor.shape[-2]:
+            raise RuntimeError(
+                "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                    self.shape, tensor.shape
+                )
+            )
 
         func = Matmul(self.representation_tree())
         return func(tensor, *self.representation())
 
+    @property
+    def matrix_shape(self):
+        """
+        Returns the shape of the matrix being represented (without batching).
+        """
+        return torch.Size(self.shape[-2:])
+
     def mul(self, other):
         """
         Multiplies the matrix by a constant, or elementwise the matrix by another matrix
@@ -791,6 +804,12 @@ class LazyTensor(object):
         """
         return len(self.size())
 
+    def numel(self):
+        """
+        Returns the number of elements
+        """
+        return self.shape.numel()
+
     def representation(self):
         """
         Returns the Tensors that are used to define the LazyTensor
@@ -820,14 +839,19 @@ class LazyTensor(object):
         This can be used for sampling from a Gaussian distribution, or for obtaining a
         low-rank version of a matrix
         """
-        batch_size = self.size(0) if self.ndimension() == 3 else None
+        if not self.is_square:
+            raise RuntimeError(
+                "root_decomposition only operates on (batches of) square (symmetric) LazyTensors. "
+                "Got a {} of size {}.".format(self.__class__.__name__, self.size())
+            )
+
         res, _ = RootDecomposition(
             self.representation_tree(),
+            max_iter=self.root_decomposition_size(),
             dtype=self.dtype,
             device=self.device,
-            size=self.size(-1),
-            max_iter=self.root_decomposition_size(),
-            batch_size=batch_size,
+            batch_shape=self.batch_shape,
+            matrix_shape=self.matrix_shape,
         )(*self.representation())
         return res
 
@@ -837,53 +861,64 @@ class LazyTensor(object):
         This can be used for sampling from a Gaussian distribution, or for obtaining a
         low-rank version of a matrix
         """
+        if not self.is_square:
+            raise RuntimeError(
+                "root_inv_decomposition only operates on (batches of) square (symmetric) LazyTensors. "
+                "Got a {} of size {}.".format(self.__class__.__name__, self.size())
+            )
+
         if initial_vectors is not None:
-            if self.ndimension() == 3:
-                if initial_vectors.ndimension() == 2:
-                    initial_vectors = initial_vectors.unsqueeze(-1)
-            else:
-                if initial_vectors.ndimension() == 1:
-                    initial_vectors = initial_vectors.unsqueeze(-1)
-            if initial_vectors.size(-1) > 1 and test_vectors is None:
-                raise RuntimeError("You must supply test vectors if you supply more than one " "initial vector")
+            if (self.dim() == 2 and initial_vectors.dim() == 1):
+                if self.shape[-1] != initial_vectors.numel():
+                    raise RuntimeError(
+                        "LazyTensor (size={}) cannot be multiplied with initial_vectors (size={}).".format(
+                            self.shape, initial_vectors.shape
+                        )
+                    )
+            elif self.dim() != initial_vectors.dim():
+                raise RuntimeError(
+                    "LazyTensor (size={}) and initial_vectors (size={}) should have the same number "
+                    "of dimensions.".format(self.shape, initial_vectors.shape)
+                )
+            elif self.batch_shape != initial_vectors.shape[:-2] or self.shape[-1] != initial_vectors.shape[-2]:
+                raise RuntimeError(
+                    "LazyTensor (size={}) cannot be multiplied with initial_vectors (size={}).".format(
+                        self.shape, initial_vectors.shape
+                    )
+                )
 
-        batch_size = self.size(0) if self.ndimension() == 3 else None
         roots, inv_roots = RootDecomposition(
             self.representation_tree(),
+            max_iter=self.root_decomposition_size(),
             dtype=self.dtype,
             device=self.device,
-            size=self.size(-1),
-            max_iter=self.root_decomposition_size(),
+            batch_shape=self.batch_shape,
+            matrix_shape=self.matrix_shape,
             root=True,
             inverse=True,
-            batch_size=batch_size,
+            initial_vectors=initial_vectors,
         )(*self.representation())
 
         # Choose the best of the inv_roots, if there were more than one initial vectors
         if initial_vectors is not None and initial_vectors.size(-1) > 1:
-            n_probes = initial_vectors.size(-1)
-            n_dim = self.size(-1)
+            num_probes = initial_vectors.size(-1)
             test_vectors = test_vectors.unsqueeze(0)
 
             # Compute solves
             solves = inv_roots.matmul(inv_roots.transpose(-1, -2).matmul(test_vectors))
 
             # Compute self * solves
-            if batch_size is None:
-                solves = solves.permute(1, 2, 0).contiguous().view(n_dim, -1)
-                mat_times_solves = self.matmul(solves)
-                mat_times_solves = mat_times_solves.view(n_dim, -1, n_probes).permute(2, 0, 1)
-            else:
-                solves = solves.permute(1, 2, 3, 0).contiguous().view(batch_size, n_dim, -1)
-                mat_times_solves = self.matmul(solves)
-                mat_times_solves = mat_times_solves.view(batch_size, n_dim, -1, n_probes).permute(3, 0, 1, 2)
+            solves = solves.permute(*range(1, self.dim() + 1), 0).contiguous().view(
+                *self.batch_shape, self.matrix_shape[-1], -1
+            )
+            mat_times_solves = self.matmul(solves)
+            mat_times_solves = mat_times_solves.view(
+                *self.batch_shape, self.matrix_shape[-1], -1, num_probes
+            ).permute(-1, *range(0, self.dim()))
 
             # Compute residuals
             residuals = (mat_times_solves - test_vectors).norm(2, dim=-2)
-            if batch_size is None:
-                residuals = residuals.sum(-1)
-            else:
-                residuals = residuals.sum(-1).sum(-1)
+            residuals = residuals.view(residuals.size(0), -1).sum(-1)
 
             # Choose solve that best fits
             _, best_solve_index = residuals.min(0)
@@ -948,6 +983,38 @@ class LazyTensor(object):
 
         return SumBatchLazyTensor(self, num_blocks=sum_batch_size)
 
+    def to(self, device_id):
+        """
+        A device-agnostic method of moving the lazy_tensor to the specified device.
+
+        Args:
+            device_id (:obj: `torch.device`): Which device to use (GPU or CPU).
+        Returns:
+            :obj:`~gpytorch.lazy.LazyTensor`: New LazyTensor identical to self on specified device
+        """
+        new_args = []
+        new_kwargs = {}
+        for arg in self._args:
+            if hasattr(arg, "to"):
+                new_args.append(arg.to(device_id))
+            else:
+                new_args.append(arg)
+        for name, val in self._kwargs.items():
+            if hasattr(val, "to"):
+                new_kwargs[name] = val.to(device_id)
+            else:
+                new_kwargs[name] = val
+        return self.__class__(*new_args, **new_kwargs)
+
+    def t(self):
+        """
+        Alias of :meth:`~gpytorch.lazy.LazyTensor.transpose` for 2D LazyTensor.
+        (Tranposes the two dimensions.)
+        """
+        if self.ndimension() != 2:
+            raise RuntimeError("Cannot call t for more than 2 dimensions")
+        return self.transpose(0, 1)
+
     def transpose(self, dim1, dim2):
         """
         Transpose the dimensions `dim1` and `dim2` of the LazyTensor.
@@ -976,15 +1043,6 @@ class LazyTensor(object):
 
         return res
 
-    def t(self):
-        """
-        Alias of :meth:`~gpytorch.lazy.LazyTensor.transpose` for 2D LazyTensor.
-        (Tranposes the two dimensions.)
-        """
-        if self.ndimension() != 2:
-            raise RuntimeError("Cannot call t for more than 2 dimensions")
-        return self.transpose(0, 1)
-
     def zero_mean_mvn_samples(self, num_samples):
         """
         Assumes that self is a covariance matrix, or a batch of covariance matrices.
@@ -1057,18 +1115,6 @@ class LazyTensor(object):
 
         return self.mul(1. / other)
 
-    def __mul__(self, other):
-        """
-        Convenience alias of :meth:`~gpytorch.lazy.LazyTensor.mul` that allows the standard product operator to be
-        used.
-        """
-        from .zero_lazy_tensor import ZeroLazyTensor
-
-        if isinstance(other, ZeroLazyTensor):
-            return other
-
-        return self.mul(other)
-
     def __getitem__(self, index):
         """
         Supports subindexing of the matrix this LazyTensor represents. This may return either another
@@ -1165,6 +1211,18 @@ class LazyTensor(object):
 
         return res
 
+    def __mul__(self, other):
+        """
+        Convenience alias of :meth:`~gpytorch.lazy.LazyTensor.mul` that allows the standard product operator to be
+        used.
+        """
+        from .zero_lazy_tensor import ZeroLazyTensor
+
+        if isinstance(other, ZeroLazyTensor):
+            return other
+
+        return self.mul(other)
+
     def __setattr__(self, name, val):
         if torch.is_tensor(val) or isinstance(val, LazyTensor):
             if not hasattr(self, "_args"):
diff --git a/gpytorch/lazy/non_lazy_tensor.py b/gpytorch/lazy/non_lazy_tensor.py
index 034e9ff3..4eb9f646 100644
--- a/gpytorch/lazy/non_lazy_tensor.py
+++ b/gpytorch/lazy/non_lazy_tensor.py
@@ -28,10 +28,6 @@ class NonLazyTensor(LazyTensor):
         return torch.matmul(self.tensor.transpose(-1, -2), rhs)
 
     def _quad_form_derivative(self, left_vecs, right_vecs):
-        if left_vecs.ndimension() < self.tensor.ndimension():
-            left_vecs = left_vecs.unsqueeze(-1)
-            right_vecs = right_vecs.unsqueeze(-1)
-
         res = left_vecs.matmul(right_vecs.transpose(-1, -2))
         return (res,)
 
@@ -55,12 +51,8 @@ class NonLazyTensor(LazyTensor):
         if self.tensor.ndimension() < 3:
             return self.tensor.diag()
         else:
-            shape = self.shape
-            batch_iter = torch.arange(0, shape[0], dtype=torch.long, device=self.device)
-            row_col_iter = torch.arange(0, shape[-1], dtype=torch.long, device=self.device)
-            batch_iter = batch_iter.unsqueeze(1).repeat(1, shape[1]).view(-1)
-            row_col_iter = row_col_iter.unsqueeze(1).repeat(shape[0], 1).view(-1)
-            return self.tensor[batch_iter, row_col_iter, row_col_iter].view(shape[:-1])
+            row_col_iter = torch.arange(0, self.matrix_shape[-1], dtype=torch.long, device=self.device)
+            return self.tensor[..., row_col_iter, row_col_iter].view(*self.batch_shape, -1)
 
     def evaluate(self):
         return self.tensor
diff --git a/gpytorch/lazy/root_lazy_tensor.py b/gpytorch/lazy/root_lazy_tensor.py
index 156fdd73..fc2fd433 100644
--- a/gpytorch/lazy/root_lazy_tensor.py
+++ b/gpytorch/lazy/root_lazy_tensor.py
@@ -37,9 +37,6 @@ class RootLazyTensor(LazyTensor):
         return self._matmul(rhs)
 
     def _quad_form_derivative(self, left_vecs, right_vecs):
-        if left_vecs.ndimension() == 1:
-            left_vecs = left_vecs.unsqueeze(0)
-            right_vecs = right_vecs.unsqueeze(0)
         right_vecs_times_rhs = self.root._t_matmul(right_vecs)
         left_vecs_times_lhs_t = self.root._t_matmul(left_vecs)
 
diff --git a/gpytorch/models/grid_inducing_variational_gp.py b/gpytorch/models/grid_inducing_variational_gp.py
index 67af6057..25f240bc 100644
--- a/gpytorch/models/grid_inducing_variational_gp.py
+++ b/gpytorch/models/grid_inducing_variational_gp.py
@@ -60,22 +60,18 @@ class GridInducingVariationalGP(AbstractVariationalGP):
 
     def __call__(self, inputs, **kwargs):
         # Prior output
-        prior_output = None
         if self.training or beta_features.diagonal_correction.on():
             prior_output = self.prior_output()
-
-        # Initialize variational parameters, if necessary
-        if not self.variational_params_initialized:
-            prior_output = prior_output if prior_output is not None else self.prior_output()
-            self._initalize_variational_parameters(prior_output)
-            self.variational_params_initialized.fill_(1)
+            if not self.variational_params_initialized.item():
+                self._initalize_variational_parameters(prior_output)
+                self.variational_params_initialized.fill_(1)
 
         # Variational output
         variational_output = self.variational_output()
 
         # Update the variational distribution
-        if self.training or self._variational_strategies["inducing_point_strategy"] is None:
-            prior_output = prior_output if prior_output is not None else self.prior_output()
+        if self.training:
+            # Initialize variational parameters, if necessary
             new_variational_strategy = MVNVariationalStrategy(variational_output, prior_output)
             self.update_variational_strategy("inducing_point_strategy", new_variational_strategy)
 
diff --git a/gpytorch/models/variational_gp.py b/gpytorch/models/variational_gp.py
index 4df0579a..71bf20ba 100644
--- a/gpytorch/models/variational_gp.py
+++ b/gpytorch/models/variational_gp.py
@@ -28,32 +28,29 @@ class VariationalGP(AbstractVariationalGP):
         return super(VariationalGP, self).train(mode)
 
     def __call__(self, inputs, **kwargs):
-        prior_output = None
-        variational_output = self.variational_output()
-
-        # Initialize variational parameters, if necessary
-        if not self.variational_params_initialized.item():
-            prior_output = self.prior_output() if prior_output is None else prior_output
-            mean_init = prior_output.mean
-            chol_covar_init = torch.eye(len(mean_init)).type_as(mean_init)
-            self.variational_mean.data.copy_(mean_init)
-            self.chol_variational_covar.data.copy_(chol_covar_init)
-            self.variational_params_initialized.fill_(1)
-
-        # Update the KL divergence term, if necessary
-        if self.training or self._variational_strategies["inducing_point_strategy"] is None:
-            prior_output = self.prior_output() if prior_output is None else prior_output
-            new_variational_strategy = MVNVariationalStrategy(variational_output, prior_output)
-            self.update_variational_strategy("inducing_point_strategy", new_variational_strategy)
-
         # Training mode: optimizing
         if self.training:
             if not torch.equal(inputs, self.inducing_points):
                 raise RuntimeError("You must train on the training inputs!")
+
+            prior_output = self.prior_output()
+            # Initialize variational parameters, if necessary
+            if not self.variational_params_initialized.item():
+                mean_init = prior_output.mean
+                chol_covar_init = torch.eye(len(mean_init)).type_as(mean_init)
+                self.variational_mean.data.copy_(mean_init)
+                self.chol_variational_covar.data.copy_(chol_covar_init)
+                self.variational_params_initialized.fill_(1)
+
+            variational_output = self.variational_output()
+            new_variational_strategy = MVNVariationalStrategy(variational_output, prior_output)
+            self.update_variational_strategy("inducing_point_strategy", new_variational_strategy)
             return variational_output
 
         # Posterior mode
         else:
+            variational_output = self.variational_output()
+
             n_induc = len(self.inducing_points)
             full_inputs = torch.cat([self.inducing_points, inputs])
             full_output = super(VariationalGP, self).__call__(full_inputs)
diff --git a/gpytorch/settings.py b/gpytorch/settings.py
index 29a62c30..27063b4a 100644
--- a/gpytorch/settings.py
+++ b/gpytorch/settings.py
@@ -134,6 +134,14 @@ class num_trace_samples(_value_context):
     _global_value = 10
 
 
+class terminate_cg_by_size(_feature_flag):
+    """
+    If set to true, cg will terminate after n iterations for an n x n matrix.
+    """
+
+    _state = True
+
+
 class use_toeplitz(_feature_flag):
     """
     Whether or not to use Toeplitz math with gridded data, grid inducing point modules
diff --git a/gpytorch/utils/cholesky.py b/gpytorch/utils/cholesky.py
index ffdb06d4..7766711d 100644
--- a/gpytorch/utils/cholesky.py
+++ b/gpytorch/utils/cholesky.py
@@ -9,20 +9,24 @@ import torch
 
 def batch_potrf(mat):
     """
+    TODO: Replace with torch batch potrf once it is implemented.
     """
-    potrf_list = []
-    for i in range(mat.size(0)):
-        potrf_list.append(mat[i].potrf().unsqueeze(0))
-    return torch.cat(potrf_list, 0)
+    potrf_list = [sub_mat.potrf() for sub_mat in mat.view(-1, *mat.shape[-2:])]
+    res = torch.cat(potrf_list, 0)
+    return res.view_as(mat)
 
 
 def batch_potrs(mat, chol):
     """
+    TODO: Replace with torch batch potrs once it is implemented.
     """
     potrs_list = []
-    for i in range(mat.size(0)):
-        potrs_list.append(torch.potrs(mat[i], chol[i]).unsqueeze(0))
-    return torch.cat(potrs_list, 0)
+    potrs_list = [
+        torch.potrs(sub_mat, sub_chol)
+        for sub_mat, sub_chol in zip(mat.view(-1, *mat.shape[-2:]), chol.view(-1, *chol.shape[-2:]))
+    ]
+    res = torch.cat(potrs_list, 0)
+    return res.view_as(mat)
 
 
 def tridiag_batch_potrf(trid, upper=False):
diff --git a/gpytorch/utils/eig.py b/gpytorch/utils/eig.py
index c00581f8..a09d9c09 100644
--- a/gpytorch/utils/eig.py
+++ b/gpytorch/utils/eig.py
@@ -10,23 +10,21 @@ def batch_symeig(mat):
     """
     """
     mat_orig = mat
+    batch_shape = torch.Size(mat_orig.shape[:-2])
+    matrix_shape = torch.Size(mat_orig.shape[-2:])
+
+    # Smaller matrices are faster on the CPU than the GPU
     if mat.size(-1) <= 32:
         mat = mat.cpu()
 
-    if mat.dim() == 3:
-        mat = mat.unsqueeze(0)
-    batch_dim1 = mat.size(0)
-    batch_dim2 = mat.size(1)
-    n = mat.size(2)
-
-    eigenvectors = torch.empty_like(mat)
-    eigenvalues = torch.empty(batch_dim1, batch_dim2, n, dtype=mat.dtype, device=mat.device)
+    mat = mat.view(-1, *matrix_shape)
+    eigenvectors = torch.empty(batch_shape.numel(), *matrix_shape, dtype=mat.dtype, device=mat.device)
+    eigenvalues = torch.empty(batch_shape.numel(), matrix_shape[-1], dtype=mat.dtype, device=mat.device)
 
-    for i in range(batch_dim1):
-        for j in range(batch_dim2):
-            evals, evecs = mat[i, j].symeig(eigenvectors=True)
-            mask = evals.ge(0)
-            eigenvectors[i, j] = evecs * mask.type_as(evecs).unsqueeze(0)
-            eigenvalues[i, j] = evals.masked_fill_(1 - mask, 1)
+    for i in range(batch_shape.numel()):
+        evals, evecs = mat[i].symeig(eigenvectors=True)
+        mask = evals.ge(0)
+        eigenvectors[i] = evecs * mask.type_as(evecs).unsqueeze(0)
+        eigenvalues[i] = evals.masked_fill_(1 - mask, 1)
 
-    return eigenvalues.type_as(mat_orig), eigenvectors.type_as(mat_orig)
+    return eigenvalues.type_as(mat_orig).view(*batch_shape, -1), eigenvectors.type_as(mat_orig).view_as(mat_orig)
diff --git a/gpytorch/utils/lanczos.py b/gpytorch/utils/lanczos.py
index 30f95bae..14b8b52d 100644
--- a/gpytorch/utils/lanczos.py
+++ b/gpytorch/utils/lanczos.py
@@ -5,79 +5,67 @@ from __future__ import unicode_literals
 
 import torch
 from .eig import batch_symeig
+from .. import settings
 
 
 def lanczos_tridiag(
     matmul_closure,
     max_iter,
-    tol=1e-5,
+    dtype,
+    device,
+    matrix_shape,
+    batch_shape=torch.Size(),
     init_vecs=None,
-    dtype=None,
-    device=None,
-    batch_size=None,
-    n_dims=None,
-    n_init_vecs=None,
+    num_init_vecs=1,
+    tol=1e-5,
 ):
     """
     """
     # Determine batch mode
-    is_batch = False
     multiple_init_vecs = False
-    if torch.is_tensor(matmul_closure):
-        lhs = matmul_closure
 
-        def default_matmul_closure(tensor):
-            return lhs.matmul(tensor)
+    if not callable(matmul_closure):
+        raise RuntimeError(
+            "matmul_closure should be a function callable object that multiples a (Lazy)Tensor "
+            "by a vector. Got a {} instead.".format(matmul_closure.__class__.__name__)
+        )
 
-        matmul_closure = default_matmul_closure
     # Get initial probe ectors - and define if not available
     if init_vecs is None:
-        if dtype is None or device is None:
-            raise RuntimeError("You must supply dtype and device if you don't supply init_vecs")
-        if n_dims is None:
-            raise RuntimeError("You must supply n_dims if you don't supply init_vecs")
-
-        if batch_size is not None:
-            is_batch = True
-        else:
-            batch_size = 1
-        if n_init_vecs is not None:
-            multiple_init_vecs = True
-        else:
-            n_init_vecs = 1
-
-        init_vecs = torch.randn(1, n_dims, n_init_vecs, dtype=dtype, device=device)
-        init_vecs = init_vecs.expand(batch_size, n_dims, n_init_vecs)
+        init_vecs = torch.randn(matrix_shape[-1], num_init_vecs, dtype=dtype, device=device)
+        init_vecs = init_vecs.expand(*batch_shape, matrix_shape[-1], num_init_vecs)
+
     else:
-        dtype = init_vecs.dtype
-        device = init_vecs.device
-        if n_init_vecs is not None or init_vecs.size(-1) > 1:
-            multiple_init_vecs = True
-        if init_vecs.ndimension() == 3:
-            batch_size, n_dims, n_init_vecs = init_vecs.size()
-            is_batch = True
-        else:
-            init_vecs = init_vecs.unsqueeze(0)
-            batch_size, n_dims, n_init_vecs = init_vecs.size()
-
-    # Modify matmul closure so that it handles batch vs non-batch
-    def batch_aware_matmul_closure(tensor):
-        if is_batch:
-            return matmul_closure(tensor)
-        else:
-            return matmul_closure(tensor.squeeze(0)).unsqueeze_(0)
+        if settings.debug.on():
+            if dtype != init_vecs.dtype:
+                raise RuntimeError(
+                    "Supplied dtype {} and init_vecs.dtype {} do not agree!".format(dtype, init_vecs.dtype)
+                )
+            if device != init_vecs.device:
+                raise RuntimeError(
+                    "Supplied device {} and init_vecs.device {} do not agree!".format(device, init_vecs.device)
+                )
+            if batch_shape != init_vecs.shape[:-2]:
+                raise RuntimeError(
+                    "batch_shape {} and init_vecs.shape {} do not agree!".format(batch_shape, init_vecs.shape)
+                )
+            if matrix_shape[-1] != init_vecs.size(-2):
+                raise RuntimeError(
+                    "matrix_shape {} and init_vecs.shape {} do not agree!".format(matrix_shape, init_vecs.shape)
+                )
+
+        num_init_vecs = init_vecs.size(-1)
 
     # Define some constants
-    n_iter = min(max_iter, n_dims)
-    batch_dimension = 1
+    num_iter = min(max_iter, matrix_shape[-1])
     dim_dimension = -2
 
     # Create storage for q_mat, alpha,and beta
     # q_mat - batch version of Q - orthogonal matrix of decomp
     # alpha - batch version main diagonal of T
     # beta - batch version of off diagonal of T
-    q_mat = torch.zeros(n_iter, batch_size, n_dims, n_init_vecs, dtype=dtype, device=device)
-    t_mat = torch.zeros(n_iter, n_iter, batch_size, n_init_vecs, dtype=dtype, device=device)
+    q_mat = torch.zeros(num_iter, *batch_shape, matrix_shape[-1], num_init_vecs, dtype=dtype, device=device)
+    t_mat = torch.zeros(num_iter, num_iter, *batch_shape, num_init_vecs, dtype=dtype, device=device)
 
     # Begin algorithm
     # Initial Q vector: q_0_vec
@@ -85,7 +73,7 @@ def lanczos_tridiag(
     q_mat[0].copy_(q_0_vec)
 
     # Initial alpha value: alpha_0
-    r_vec = batch_aware_matmul_closure(q_0_vec)
+    r_vec = matmul_closure(q_0_vec)
     alpha_0 = q_0_vec.mul(r_vec).sum(dim_dimension)
 
     # Initial beta value: beta_0
@@ -101,20 +89,20 @@ def lanczos_tridiag(
     q_mat[1].copy_(r_vec.div_(beta_0.unsqueeze(dim_dimension)))
 
     # Now we start the iteration
-    for k in range(1, n_iter):
+    for k in range(1, num_iter):
         # Get previous values
         q_prev_vec = q_mat[k - 1]
         q_curr_vec = q_mat[k]
         beta_prev = t_mat[k, k - 1].unsqueeze(dim_dimension)
 
         # Compute next alpha value
-        r_vec = batch_aware_matmul_closure(q_curr_vec) - q_prev_vec.mul(beta_prev)
+        r_vec = matmul_closure(q_curr_vec) - q_prev_vec.mul(beta_prev)
         alpha_curr = q_curr_vec.mul(r_vec).sum(dim_dimension, keepdim=True)
         # Copy over to t_mat
         t_mat[k, k].copy_(alpha_curr.squeeze(dim_dimension))
 
         # Copy over alpha_curr, beta_curr to t_mat
-        if (k + 1) < n_iter:
+        if (k + 1) < num_iter:
             # Compute next residual value
             r_vec.sub_(alpha_curr.mul(q_curr_vec))
             # Full reorthogonalization: r <- r - Q (Q^T r)
@@ -151,17 +139,14 @@ def lanczos_tridiag(
                 break
 
     # Now let's transpose q_mat, t_mat intot the correct shape
-    n_iter = k + 1
+    num_iter = k + 1
 
-    # n_init_vecs x batch_size x n_dims x n_iter
-    q_mat = q_mat[: n_iter + 1].permute(3, 1, 2, 0).contiguous()
-    # n_init_vecs x batch_size x n_iter x n_iter
-    t_mat = t_mat[: n_iter + 1, : n_iter + 1].permute(3, 2, 0, 1).contiguous()
+    # num_init_vecs x batch_shape x matrix_shape[-1] x num_iter
+    q_mat = q_mat[: num_iter + 1].permute(-1, *range(1, 1 + len(batch_shape)), -2, 0).contiguous()
+    # num_init_vecs x batch_shape x num_iter x num_iter
+    t_mat = t_mat[: num_iter + 1, : num_iter + 1].permute(-1, *range(2, 2 + len(batch_shape)), 0, 1).contiguous()
 
     # If we weren't in batch mode, remove batch dimension
-    if not is_batch:
-        q_mat.squeeze_(batch_dimension)
-        t_mat.squeeze_(batch_dimension)
     if not multiple_init_vecs:
         q_mat.squeeze_(0)
         t_mat.squeeze_(0)
diff --git a/gpytorch/utils/linear_cg.py b/gpytorch/utils/linear_cg.py
index e138a8cc..3ed86814 100644
--- a/gpytorch/utils/linear_cg.py
+++ b/gpytorch/utils/linear_cg.py
@@ -70,8 +70,9 @@ def linear_cg(
         raise RuntimeError("matmul_closure must be a tensor, or a callable object!")
 
     # Get some constants
+    batch_shape = rhs.shape[:-2]
     num_rows = rhs.size(-2)
-    n_iter = min(max_iter, num_rows)
+    n_iter = min(max_iter, num_rows) if settings.terminate_cg_by_size.on() else max_iter
     n_tridiag_iter = min(max_tridiag_iter, num_rows)
 
     # result <- x_{0}
@@ -98,23 +99,15 @@ def linear_cg(
 
         # Define storage matrices
         mul_storage = torch.empty_like(residual)
-        if rhs.ndimension() == 3:
-            alpha = torch.empty(rhs.size(0), 1, rhs.size(-1), dtype=residual.dtype, device=residual.device)
-        else:
-            alpha = torch.empty(1, rhs.size(-1), dtype=residual.dtype, device=residual.device)
+        alpha = torch.empty(*batch_shape, rhs.size(-1), dtype=residual.dtype, device=residual.device)
         beta = torch.empty_like(alpha)
 
     # Define tridiagonal matrices, if applicable
     if n_tridiag:
-        if rhs.ndimension() == 3:
-            t_mat = torch.zeros(
-                n_tridiag_iter, n_tridiag_iter, rhs.size(0), n_tridiag, dtype=alpha.dtype, device=alpha.device
-            )
-            alpha_reciprocal = torch.empty(rhs.size(0), n_tridiag, dtype=t_mat.dtype, device=t_mat.device)
-        else:
-            t_mat = torch.zeros(n_tridiag_iter, n_tridiag_iter, n_tridiag, dtype=alpha.dtype, device=alpha.device)
-            alpha_reciprocal = torch.empty(n_tridiag, dtype=t_mat.dtype, device=t_mat.device)
-
+        t_mat = torch.zeros(
+            n_tridiag_iter, n_tridiag_iter, *batch_shape, n_tridiag, dtype=alpha.dtype, device=alpha.device
+        )
+        alpha_reciprocal = torch.empty(*batch_shape, n_tridiag, dtype=t_mat.dtype, device=t_mat.device)
         prev_alpha_reciprocal = torch.empty_like(alpha_reciprocal)
         prev_beta = torch.empty_like(alpha_reciprocal)
 
@@ -185,9 +178,6 @@ def linear_cg(
 
     if n_tridiag:
         t_mat = t_mat[: last_tridiag_iter + 1, : last_tridiag_iter + 1]
-        if rhs.ndimension() == 3:
-            return result, t_mat.permute(3, 2, 0, 1).contiguous()
-        else:
-            return result, t_mat.permute(2, 0, 1).contiguous()
+        return result, t_mat.permute(-1, *range(2, 2 + len(batch_shape)), 0, 1).contiguous()
     else:
         return result
diff --git a/gpytorch/utils/pivoted_cholesky.py b/gpytorch/utils/pivoted_cholesky.py
index b7e51a21..0f01485e 100644
--- a/gpytorch/utils/pivoted_cholesky.py
+++ b/gpytorch/utils/pivoted_cholesky.py
@@ -5,14 +5,8 @@ from .cholesky import batch_potrf, batch_potrs
 def pivoted_cholesky(matrix, max_iter, error_tol=1e-3):
     from ..lazy import LazyTensor, NonLazyTensor
 
-    # matrix is assumed to be batch_size x n x n
-    if matrix.ndimension() < 3:
-        batch_size = 1
-        batch_mode = False
-    else:
-        batch_size = matrix.size(0)
-        batch_mode = True
-    matrix_size = matrix.size(-1)
+    batch_shape = matrix.shape[:-2]
+    matrix_shape = matrix.shape[-2:]
 
     # Need to get diagonals. This is easy if it's a LazyTensor, since
     # LazyTensor.diag() operates in batch mode.
@@ -22,70 +16,67 @@ def pivoted_cholesky(matrix, max_iter, error_tol=1e-3):
     elif torch.is_tensor(matrix):
         matrix_diag = NonLazyTensor(matrix).diag()
 
-    if not batch_mode:
-        matrix_diag.unsqueeze_(0)
-    # matrix_diag is now batch_size x n
-
     # Make sure max_iter isn't bigger than the matrix
-    max_iter = min(max_iter, matrix_size)
+    max_iter = min(max_iter, matrix_shape[-1])
+
+    # What we're returning
+    L = torch.zeros(*batch_shape, max_iter, matrix_shape[-1], dtype=matrix.dtype, device=matrix.device)
+    errors = torch.norm(matrix_diag, 1, dim=-1)
+
+    # The permutation
+    permutation = torch.arange(0, matrix_shape[-1], dtype=torch.long, device=matrix_diag.device)
+    permutation = permutation.repeat(*batch_shape, 1)
 
-    errors = torch.norm(matrix_diag, 1, dim=1)
-    permutation = torch.arange(0, matrix_size, dtype=torch.long, device=matrix_diag.device)
-    permutation = permutation.repeat(batch_size, 1)
+    # Get batch indices
+    batch_iters = [
+        torch.arange(0, size, dtype=torch.long, device=matrix_diag.device).unsqueeze_(-1).repeat(
+            torch.Size(batch_shape[:i]).numel(), torch.Size(batch_shape[i + 1:]).numel()
+        ).view(-1)
+        for i, size in enumerate(batch_shape)
+    ]
 
     m = 0
-    L = torch.zeros(batch_size, max_iter, matrix_size, dtype=matrix.dtype, device=matrix.device)
-    full_batch_slice = torch.arange(0, batch_size, dtype=torch.long, device=permutation.device)
     while m < max_iter and torch.max(errors) > error_tol:
-        permuted_diags = torch.gather(matrix_diag, 1, permutation)[:, m:]
-        max_diag_values, max_diag_indices = torch.max(permuted_diags, 1)
+        permuted_diags = torch.gather(matrix_diag, -1, permutation[..., m:])
+        max_diag_values, max_diag_indices = torch.max(permuted_diags, -1)
 
         max_diag_indices = max_diag_indices + m
 
         # Swap pi_m and pi_i in each row, where pi_i is the element of the permutation
         # corresponding to the max diagonal element
-        old_pi_m = permutation[:, m].clone()
-        new_pi_m = permutation[full_batch_slice, max_diag_indices].clone()
-        permutation[:, m] = new_pi_m
-        permutation[full_batch_slice, max_diag_indices] = old_pi_m
-        pi_m = permutation[:, m]
-
-        L_m = L[:, m]  # Will be all zeros -- should we use torch.zeros?
-        L_m[full_batch_slice, pi_m] = torch.sqrt(max_diag_values)
-
-        if not batch_mode:
-            row = matrix[pi_m, :]
-            if row.ndimension() < 2:
-                row.unsqueeze_(0)
-        else:
-            row = matrix[full_batch_slice, pi_m, :]
+        old_pi_m = permutation[..., m].clone()
+        permutation[..., m].copy_(permutation.gather(-1, max_diag_indices.unsqueeze(-1)).squeeze_(-1))
+        permutation.scatter_(-1, max_diag_indices.unsqueeze(-1), old_pi_m.unsqueeze(-1))
+        pi_m = permutation[..., m].contiguous()
 
+        L_m = L[..., m, :]  # Will be all zeros -- should we use torch.zeros?
+        L_m.scatter_(-1, pi_m.unsqueeze(-1), max_diag_values.sqrt().unsqueeze_(-1))
+
+        row = matrix[(*batch_iters, pi_m.view(-1), slice(None, None, None))]
         if isinstance(row, LazyTensor):
             row = row.evaluate()
+        row = row.view(*batch_shape, matrix_shape[-1])
 
-        if m + 1 < matrix_size:
-            pi_i = permutation[:, m + 1 :]
+        if m + 1 < matrix_shape[-1]:
+            pi_i = permutation[..., m + 1:].contiguous()
 
-            L_m_new = row.gather(1, pi_i)
+            L_m_new = row.gather(-1, pi_i)
             if m > 0:
-                L_prev = L[:, :m].gather(2, pi_i.unsqueeze(1).repeat(1, m, 1))
-                update = L[:, :m].gather(2, pi_m.unsqueeze(1).unsqueeze(1).repeat(1, m, 1))
-                L_m_new -= torch.sum(update * L_prev, dim=1)
+                L_prev = L[..., :m, :].gather(-1, pi_i.unsqueeze(-2).repeat(*(1 for _ in batch_shape), m, 1))
+                update = L[..., :m, :].gather(-1, pi_m.view(*pi_m.shape, 1, 1).repeat(*(1 for _ in batch_shape), m, 1))
+                L_m_new -= torch.sum(update * L_prev, dim=-2)
 
-            L_m_new /= L_m.gather(1, pi_m.unsqueeze(1))
-            L_m.scatter_(1, pi_i, L_m_new)
+            L_m_new /= L_m.gather(-1, pi_m.unsqueeze(-1))
+            L_m.scatter_(-1, pi_i, L_m_new)
 
-            matrix_diag_current = matrix_diag.gather(1, pi_i)
-            matrix_diag.scatter_(1, pi_i, matrix_diag_current - L_m_new ** 2)
-            L[:, m] = L_m
+            matrix_diag_current = matrix_diag.gather(-1, pi_i)
+            matrix_diag.scatter_(-1, pi_i, matrix_diag_current - L_m_new ** 2)
+            L[..., m, :] = L_m
 
-            errors = torch.norm(matrix_diag.gather(1, pi_i), 1, dim=1)
+            errors = torch.norm(matrix_diag.gather(-1, pi_i), 1, dim=-1)
         m = m + 1
 
-    if not batch_mode:
-        return L[0, :m, :]
-    else:
-        return L[:, :m, :]
+    return L[..., :m, :].contiguous()
 
 
 def woodbury_factor(low_rank_mat, shift):
@@ -100,11 +91,7 @@ def woodbury_factor(low_rank_mat, shift):
 
     shifted_mat = shifted_mat + torch.eye(k, dtype=shifted_mat.dtype, device=shifted_mat.device)
 
-    if low_rank_mat.ndimension() == 3:
-        R = batch_potrs(low_rank_mat, batch_potrf(shifted_mat))
-    else:
-        R = torch.potrs(low_rank_mat, shifted_mat.potrf())
-
+    R = batch_potrs(low_rank_mat, batch_potrf(shifted_mat))
     return R
 
 
diff --git a/gpytorch/utils/stochastic_lq.py b/gpytorch/utils/stochastic_lq.py
index 61d15644..06d5f483 100644
--- a/gpytorch/utils/stochastic_lq.py
+++ b/gpytorch/utils/stochastic_lq.py
@@ -33,10 +33,11 @@ class StochasticLQ(object):
 
     def lanczos_batch(self, matmul_closure, rhs_vectors):
         return lanczos_tridiag(
-            matmul_closure, self.max_iter, init_vecs=rhs_vectors, dtype=rhs_vectors.dtype, device=rhs_vectors.device
+            matmul_closure, self.max_iter, init_vecs=rhs_vectors, dtype=rhs_vectors.dtype, device=rhs_vectors.device,
+            batch_shape=rhs_vectors.shape[-2:], matrix_shape=torch.Size((rhs_vectors.size(-2), rhs_vectors.size(-2)))
         )
 
-    def evaluate(self, t_mats, matrix_size, eigenvalues, eigenvectors, funcs):
+    def evaluate(self, matrix_shape, eigenvalues, eigenvectors, funcs):
         """
         Computes tr(f(A)) for an arbitrary list of functions, where f(A) is equivalent to applying the function
         elementwise to the eigenvalues of A, i.e., if A = V\LambdaV^{T}, then f(A) = Vf(\Lambda)V^{T}, where
@@ -45,10 +46,9 @@ class StochasticLQ(object):
         calling it multiple times with one function -- each additional function after the first requires negligible
         additional computation.
         Args:
-            - A (matrix n x n or closure) - Either the input matrix A or a closure that takes an n dimensional vector v
-                and returns Av.
-            - n (scalar) - dimension of the matrix A. We require this because, if A is passed in as a closure, we
-                have no good way of determining the size of A.
+            - matrix_shape (torch.Size()) - size of underlying matrix (not including batch dimensions)
+            - eigenvalues (Tensor n_probes x ...batch_shape x k) - batches of eigenvalues from Lanczos tridiag mats
+            - eigenvectors (Tensor n_probes x ...batch_shape x k x k) - batches of eigenvectors from " " "
             - funcs (list of closures) - A list of functions [f_1,...,f_k]. tr(f_i(A)) is computed for each function.
                     Each function in the closure should expect to take a torch vector of eigenvalues as input and apply
                     the function elementwise. For example, to compute logdet(A) = tr(log(A)), [lambda x: x.log()] would
@@ -57,26 +57,19 @@ class StochasticLQ(object):
             - results (list of scalars) - The trace of each supplied function applied to the matrix, e.g.,
                       [tr(f_1(A)),tr(f_2(A)),...,tr(f_k(A))].
         """
-        if t_mats.dim() < 4:
-            raise RuntimeError(
-                "StochasticLQ expects t_mat to be (num_probe_vecs x num_batch x k x k), "
-                "but received a Tensor with only {} dimensions. Try unsqueezing to add "
-                "appropriate singleton dimensions if necessary.".format(t_mats.dim())
-            )
-
-        batch_size = t_mats.size(1)
-        results = [torch.zeros(batch_size, dtype=t_mats.dtype, device=t_mats.device) for _ in funcs]
-        num_random_probes = t_mats.size(0)
+        batch_shape = torch.Size(eigenvalues.shape[1:-1])
+        results = [torch.zeros(batch_shape, dtype=eigenvalues.dtype, device=eigenvalues.device) for _ in funcs]
+        num_random_probes = eigenvalues.size(0)
         for j in range(num_random_probes):
             # These are (num_batch x k) and (num_batch x k x k)
             eigenvalues_for_probe = eigenvalues[j]
             eigenvectors_for_probe = eigenvectors[j]
             for i, func in enumerate(funcs):
                 # First component of eigenvecs is (num_batch x k)
-                eigenvecs_first_component = eigenvectors_for_probe[:, 0, :]
+                eigenvecs_first_component = eigenvectors_for_probe[..., 0, :]
                 func_eigenvalues = func(eigenvalues_for_probe)
 
-                dot_products = (eigenvecs_first_component.pow(2) * func_eigenvalues).sum(1)
-                results[i] = results[i] + matrix_size / float(num_random_probes) * dot_products
+                dot_products = (eigenvecs_first_component.pow(2) * func_eigenvalues).sum(-1)
+                results[i] = results[i] + matrix_shape[-1] / float(num_random_probes) * dot_products
 
         return results
diff --git a/requirements.txt b/requirements.txt
index 28ae7d81..b54f97d6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,1 +1,1 @@
-torch>=1.0.0
+torch>=0.4.1


Test patch:
diff --git a/test/examples/test_simple_gp_classification.py b/test/examples/test_simple_gp_classification.py
index f38ab9d9..6b86d382 100644
--- a/test/examples/test_simple_gp_classification.py
+++ b/test/examples/test_simple_gp_classification.py
@@ -81,6 +81,7 @@ class TestSimpleGPClassification(unittest.TestCase):
         for param in likelihood.parameters():
             self.assertTrue(param.grad is not None)
             self.assertGreater(param.grad.norm().item(), 0)
+        optimizer.step()
 
         # Set back to eval mode
         model.eval()
diff --git a/test/functions/test_inv_matmul.py b/test/functions/test_inv_matmul.py
index 112f89cb..3ef192be 100644
--- a/test/functions/test_inv_matmul.py
+++ b/test/functions/test_inv_matmul.py
@@ -5,74 +5,141 @@ from __future__ import unicode_literals
 
 import torch
 import unittest
+import os
+import random
+from gpytorch import settings
 from gpytorch.lazy import NonLazyTensor
-from test._utils import approx_equal
 
 
 class TestInvMatmulNonBatch(unittest.TestCase):
+    def tearDown(self):
+        if hasattr(self, "rng_state"):
+            torch.set_rng_state(self.rng_state)
+
     def setUp(self):
-        mat = [[3, -1, 0], [-1, 3, 0], [0, 0, 3]]
-        vec = torch.randn(3, requires_grad=True)
-        vecs = torch.randn(3, 4, requires_grad=True)
+        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
+            self.rng_state = torch.get_rng_state()
+            torch.manual_seed(1)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(1)
+            random.seed(1)
 
-        self.mat_var = torch.tensor(mat, dtype=torch.float, requires_grad=True)
-        self.mat_var_copy = torch.tensor(mat, dtype=torch.float, requires_grad=True)
-        self.vec_var = vec
-        self.vec_var_copy = vec.clone().detach().requires_grad_(True)
-        self.vecs_var = vecs
-        self.vecs_var_copy = vecs.clone().detach().requires_grad_(True)
+        mat = torch.randn(8, 8)
+        mat = mat @ mat.transpose(-1, -2)
+        vec = torch.randn(8)
+        vecs = torch.randn(8, 4)
+        self.mat = mat.detach().clone().requires_grad_(True)
+        self.mat_copy = mat.detach().clone().requires_grad_(True)
+        self.vec = vec.detach().clone().requires_grad_(True)
+        self.vec_copy = vec.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_copy = vecs.detach().clone().requires_grad_(True)
 
     def test_inv_matmul_vec(self):
         # Forward
-        res = NonLazyTensor(self.mat_var).inv_matmul(self.vec_var)
-        actual = self.mat_var_copy.inverse().matmul(self.vec_var_copy)
-        self.assertTrue(approx_equal(res, actual))
+        with settings.terminate_cg_by_size(False):
+            res = NonLazyTensor(self.mat).inv_matmul(self.vec)
+            actual = self.mat_copy.inverse().matmul(self.vec_copy)
+            self.assertLess(torch.max((res - actual).abs() / actual.abs()).item(), 1e-3)
 
-        # Backward
-        grad_output = torch.randn(3)
-        res.backward(gradient=grad_output)
-        actual.backward(gradient=grad_output)
-        self.assertTrue(approx_equal(self.mat_var_copy.grad, self.mat_var.grad))
-        self.assertTrue(approx_equal(self.vec_var_copy.grad, self.vec_var.grad))
+            # Backward
+            grad_output = torch.randn(8)
+            res.backward(gradient=grad_output)
+            actual.backward(gradient=grad_output)
+            self.assertLess(torch.max((self.mat_copy.grad - self.mat.grad).abs()).item(), 1e-3)
+            self.assertLess(torch.max((self.vec_copy.grad - self.vec.grad).abs()).item(), 1e-3)
 
     def test_inv_matmul_multiple_vecs(self):
         # Forward
-        res = NonLazyTensor(self.mat_var).inv_matmul(self.vecs_var)
-        actual = self.mat_var_copy.inverse().matmul(self.vecs_var_copy)
-        self.assertTrue(approx_equal(res, actual))
+        with settings.terminate_cg_by_size(False):
+            res = NonLazyTensor(self.mat).inv_matmul(self.vecs)
+            actual = self.mat_copy.inverse().matmul(self.vecs_copy)
+            self.assertLess(torch.max((res - actual).abs() / actual.abs()).item(), 1e-3)
 
-        # Backward
-        grad_output = torch.randn(3, 4)
-        res.backward(gradient=grad_output)
-        actual.backward(gradient=grad_output)
-        self.assertTrue(approx_equal(self.mat_var_copy.grad, self.mat_var.grad))
-        self.assertTrue(approx_equal(self.vecs_var_copy.grad, self.vecs_var.grad))
+            # Backward
+            grad_output = torch.randn(8, 4)
+            res.backward(gradient=grad_output)
+            actual.backward(gradient=grad_output)
+            self.assertLess(torch.max((self.mat_copy.grad - self.mat.grad).abs()).item(), 1e-3)
+            self.assertLess(torch.max((self.vecs_copy.grad - self.vecs.grad).abs()).item(), 1e-3)
 
 
 class TestInvMatmulBatch(unittest.TestCase):
+    def tearDown(self):
+        if hasattr(self, "rng_state"):
+            torch.set_rng_state(self.rng_state)
+
+    def setUp(self):
+        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
+            self.rng_state = torch.get_rng_state()
+            torch.manual_seed(0)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(0)
+            random.seed(0)
+
+        mats = torch.randn(2, 8, 8)
+        mats = mats @ mats.transpose(-1, -2)
+        vecs = torch.randn(2, 8, 4)
+        self.mats = mats.detach().clone().requires_grad_(True)
+        self.mats_copy = mats.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_copy = vecs.detach().clone().requires_grad_(True)
+
+    def test_inv_matmul_multiple_vecs(self):
+        # Forward
+        with settings.terminate_cg_by_size(False):
+            res = NonLazyTensor(self.mats).inv_matmul(self.vecs)
+            actual = torch.cat(
+                [self.mats_copy[0].inverse().unsqueeze(0), self.mats_copy[1].inverse().unsqueeze(0)]
+            ).matmul(self.vecs_copy)
+            self.assertLess(torch.max((res - actual).abs()).item(), 1e-3)
+
+            # Backward
+            grad_output = torch.randn(2, 8, 4)
+            res.backward(gradient=grad_output)
+            actual.backward(gradient=grad_output)
+            self.assertLess(torch.max((self.mats_copy.grad - self.mats.grad).abs()).item(), 1e-3)
+            self.assertLess(torch.max((self.vecs_copy.grad - self.vecs.grad).abs()).item(), 1e-3)
+
+
+class TestInvMatmulMultiBatch(unittest.TestCase):
+    def tearDown(self):
+        if hasattr(self, "rng_state"):
+            torch.set_rng_state(self.rng_state)
+
     def setUp(self):
-        mats = [[[3, -1, 0], [-1, 3, 0], [0, 0, 3]], [[10, -2, 1], [-2, 10, 0], [1, 0, 10]]]
-        vecs = torch.randn(2, 3, 4, requires_grad=True)
+        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
+            self.rng_state = torch.get_rng_state()
+            torch.manual_seed(0)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(0)
+            random.seed(0)
 
-        self.mats_var = torch.tensor(mats, dtype=torch.float, requires_grad=True)
-        self.mats_var_copy = torch.tensor(mats, dtype=torch.float, requires_grad=True)
-        self.vecs_var = vecs
-        self.vecs_var_copy = vecs.clone().detach().requires_grad_(True)
+        mats = torch.randn(3, 4, 8, 8)
+        mats = mats @ mats.transpose(-1, -2)
+        mats.add_(torch.eye(8).mul_(1e-1).view(1, 1, 8, 8))
+        vecs = torch.randn(3, 4, 8, 2)
+        self.mats = mats.detach().clone().requires_grad_(True)
+        self.mats_copy = mats.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_copy = vecs.detach().clone().requires_grad_(True)
 
     def test_inv_matmul_multiple_vecs(self):
         # Forward
-        res = NonLazyTensor(self.mats_var).inv_matmul(self.vecs_var)
-        actual = torch.cat(
-            [self.mats_var_copy[0].inverse().unsqueeze(0), self.mats_var_copy[1].inverse().unsqueeze(0)]
-        ).matmul(self.vecs_var_copy)
-        self.assertTrue(approx_equal(res, actual))
-
-        # Backward
-        grad_output = torch.randn(2, 3, 4)
-        res.backward(gradient=grad_output)
-        actual.backward(gradient=grad_output)
-        self.assertTrue(approx_equal(self.mats_var_copy.grad, self.mats_var.grad))
-        self.assertTrue(approx_equal(self.vecs_var_copy.grad, self.vecs_var.grad))
+        with settings.terminate_cg_by_size(False):
+            res = NonLazyTensor(self.mats).inv_matmul(self.vecs)
+            flattened_mats_copy = self.mats_copy.view(-1, *self.mats.shape[-2:])
+            flatened_mats_inverse = torch.cat([mat.inverse().unsqueeze(0) for mat in flattened_mats_copy])
+            mats_inverse = flatened_mats_inverse.view_as(self.mats)
+            actual = mats_inverse.matmul(self.vecs_copy)
+            self.assertLess(torch.max((res - actual).abs()).item(), 1e-3)
+
+            # Backward
+            grad_output = torch.randn(3, 4, 8, 2)
+            res.backward(gradient=grad_output)
+            actual.backward(gradient=grad_output)
+            self.assertLess(torch.max((self.mats_copy.grad - self.mats.grad).abs()).item(), 1e-3)
+            self.assertLess(torch.max((self.vecs_copy.grad - self.vecs.grad).abs()).item(), 1e-3)
 
 
 if __name__ == "__main__":
diff --git a/test/functions/test_inv_quad_log_det.py b/test/functions/test_inv_quad_log_det.py
index cd5375bb..86ac0a7a 100644
--- a/test/functions/test_inv_quad_log_det.py
+++ b/test/functions/test_inv_quad_log_det.py
@@ -4,44 +4,48 @@ from __future__ import print_function
 from __future__ import unicode_literals
 
 import os
+import random
 import torch
 import unittest
 import gpytorch
 from gpytorch.lazy import NonLazyTensor
-from test._utils import approx_equal
 
 
 class TestInvQuadLogDetNonBatch(unittest.TestCase):
-    def setUp(self):
-        if os.getenv("UNLOCK_SEED") is None or os.getenv("UNLOCK_SEED").lower() == "false":
-            self.rng_state = torch.get_rng_state()
-            torch.manual_seed(1)
-
-        self.mat_var = torch.tensor([[3, -1, 0], [-1, 3, 0], [0, 0, 3]], dtype=torch.float)
-        self.mat_var_clone = self.mat_var.clone()
-        self.vec_var = torch.randn(3)
-        self.vec_var_clone = self.vec_var.clone()
-        self.vecs_var = torch.randn(3, 4)
-        self.vecs_var_clone = self.vecs_var.clone()
-        self.mat_var.requires_grad = True
-        self.mat_var_clone.requires_grad = True
-        self.vec_var.requires_grad = True
-        self.vec_var_clone.requires_grad = True
-        self.vecs_var.requires_grad = True
-        self.vecs_var_clone.requires_grad = True
-
     def tearDown(self):
         if hasattr(self, "rng_state"):
             torch.set_rng_state(self.rng_state)
 
+    def setUp(self):
+        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
+            self.rng_state = torch.get_rng_state()
+            torch.manual_seed(0)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(0)
+            random.seed(0)
+
+        mat = torch.randn(4, 4)
+        mat = mat @ mat.transpose(-1, -2)
+        mat.div_(5).add_(torch.eye(4))
+        vecs = torch.randn(5, 4, 6)
+        vec = torch.randn(4)
+        vecs = torch.randn(4, 8)
+        self.mat = mat.detach().clone().requires_grad_(True)
+        self.mat_clone = mat.detach().clone().requires_grad_(True)
+        self.vec = vec.detach().clone().requires_grad_(True)
+        self.vec_clone = vec.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_clone = vecs.detach().clone().requires_grad_(True)
+
     def test_inv_quad_log_det_vector(self):
         # Forward pass
-        actual_inv_quad = self.mat_var_clone.inverse().matmul(self.vec_var_clone).mul(self.vec_var_clone).sum()
-        actual_log_det = self.mat_var_clone.det().log()
+        actual_inv_quad = self.mat_clone.inverse().matmul(self.vec_clone).mul(self.vec_clone).sum()
+        actual_log_det = self.mat_clone.logdet()
         with gpytorch.settings.num_trace_samples(1000):
-            nlv = NonLazyTensor(self.mat_var)
-            res_inv_quad, res_log_det = nlv.inv_quad_log_det(inv_quad_rhs=self.vec_var, log_det=True)
-        self.assertAlmostEqual(res_inv_quad, actual_inv_quad, places=1)
+            non_lazy_tsr = NonLazyTensor(self.mat)
+            res_inv_quad, res_log_det = non_lazy_tsr.inv_quad_log_det(inv_quad_rhs=self.vec, log_det=True)
+
+        self.assertAlmostEqual(res_inv_quad.item(), actual_inv_quad.item(), places=1)
         self.assertAlmostEqual(res_log_det.item(), actual_log_det.item(), places=1)
 
         # Backward
@@ -50,29 +54,29 @@ class TestInvQuadLogDetNonBatch(unittest.TestCase):
         res_inv_quad.backward(retain_graph=True)
         res_log_det.backward()
 
-        self.assertTrue(approx_equal(self.mat_var_clone.grad, self.mat_var.grad, epsilon=1e-1))
-        self.assertTrue(approx_equal(self.vec_var_clone.grad, self.vec_var.grad))
+        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((self.vec_clone.grad - self.vec.grad).abs()).item(), 1e-1)
 
     def test_inv_quad_only_vector(self):
         # Forward pass
-        res = NonLazyTensor(self.mat_var).inv_quad(self.vec_var)
-        actual = self.mat_var_clone.inverse().matmul(self.vec_var_clone).mul(self.vec_var_clone).sum()
+        res = NonLazyTensor(self.mat).inv_quad(self.vec)
+        actual = self.mat_clone.inverse().matmul(self.vec_clone).mul(self.vec_clone).sum()
         self.assertAlmostEqual(res.item(), actual.item(), places=1)
 
         # Backward
         actual.backward()
         res.backward()
 
-        self.assertTrue(approx_equal(self.mat_var_clone.grad, self.mat_var.grad, epsilon=1e-1))
-        self.assertTrue(approx_equal(self.vec_var_clone.grad, self.vec_var.grad))
+        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-3)
+        self.assertLess(torch.max((self.vec_clone.grad - self.vec.grad).abs()).item(), 1e-3)
 
     def test_inv_quad_log_det_many_vectors(self):
         # Forward pass
-        actual_inv_quad = self.mat_var_clone.inverse().matmul(self.vecs_var_clone).mul(self.vecs_var_clone).sum()
-        actual_log_det = self.mat_var_clone.det().log()
+        actual_inv_quad = self.mat_clone.inverse().matmul(self.vecs_clone).mul(self.vecs_clone).sum()
+        actual_log_det = self.mat_clone.logdet()
         with gpytorch.settings.num_trace_samples(1000):
-            nlv = NonLazyTensor(self.mat_var)
-            res_inv_quad, res_log_det = nlv.inv_quad_log_det(inv_quad_rhs=self.vecs_var, log_det=True)
+            non_lazy_tsr = NonLazyTensor(self.mat)
+            res_inv_quad, res_log_det = non_lazy_tsr.inv_quad_log_det(inv_quad_rhs=self.vecs, log_det=True)
         self.assertAlmostEqual(res_inv_quad.item(), actual_inv_quad.item(), places=1)
         self.assertAlmostEqual(res_log_det.item(), actual_log_det.item(), places=1)
 
@@ -82,112 +86,215 @@ class TestInvQuadLogDetNonBatch(unittest.TestCase):
         res_inv_quad.backward(retain_graph=True)
         res_log_det.backward()
 
-        self.assertTrue(approx_equal(self.mat_var_clone.grad, self.mat_var.grad, epsilon=1e-1))
-        self.assertTrue(approx_equal(self.vecs_var_clone.grad, self.vecs_var.grad))
+        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
 
     def test_inv_quad_only_many_vectors(self):
         # Forward pass
-        res = NonLazyTensor(self.mat_var).inv_quad(self.vecs_var)
-        actual = self.mat_var_clone.inverse().matmul(self.vecs_var_clone).mul(self.vecs_var_clone).sum()
+        res = NonLazyTensor(self.mat).inv_quad(self.vecs)
+        actual = self.mat_clone.inverse().matmul(self.vecs_clone).mul(self.vecs_clone).sum()
         self.assertAlmostEqual(res.item(), actual.item(), places=1)
 
         # Backward
         actual.backward()
         res.backward()
 
-        self.assertTrue(approx_equal(self.mat_var_clone.grad, self.mat_var.grad, epsilon=1e-1))
-        self.assertTrue(approx_equal(self.vecs_var_clone.grad, self.vecs_var.grad))
+        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-3)
+        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-3)
 
     def test_log_det_only(self):
         # Forward pass
         with gpytorch.settings.num_trace_samples(1000):
-            res = NonLazyTensor(self.mat_var).log_det()
-        actual = self.mat_var_clone.det().log()
+            res = NonLazyTensor(self.mat).log_det()
+        actual = self.mat_clone.logdet()
         self.assertAlmostEqual(res.item(), actual.item(), places=1)
 
         # Backward
         actual.backward()
         res.backward()
-        self.assertTrue(approx_equal(self.mat_var_clone.grad, self.mat_var.grad, epsilon=1e-1))
+        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-1)
 
 
 class TestInvQuadLogDetBatch(unittest.TestCase):
-    def setUp(self):
-        if os.getenv("UNLOCK_SEED") is None or os.getenv("UNLOCK_SEED").lower() == "false":
-            self.rng_state = torch.get_rng_state()
-            torch.manual_seed(1)
-
-        mats = [[[3, -1, 0], [-1, 3, 0], [0, 0, 3]], [[10, -2, 1], [-2, 10, 0], [1, 0, 10]]]
-        self.mats_var = torch.tensor(mats, dtype=torch.float, requires_grad=True)
-        self.mats_var_clone = self.mats_var.clone().detach().requires_grad_(True)
-        self.vecs_var = torch.randn(2, 3, 4, requires_grad=True)
-        self.vecs_var_clone = self.vecs_var.clone().detach().requires_grad_(True)
-
     def tearDown(self):
         if hasattr(self, "rng_state"):
             torch.set_rng_state(self.rng_state)
 
+    def setUp(self):
+        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
+            self.rng_state = torch.get_rng_state()
+            torch.manual_seed(1)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(1)
+            random.seed(1)
+
+        mats = torch.randn(5, 4, 4)
+        mats = mats @ mats.transpose(-1, -2)
+        mats.div_(5).add_(torch.eye(4).unsqueeze_(0))
+        vecs = torch.randn(5, 4, 6)
+        self.mats = mats.detach().clone().requires_grad_(True)
+        self.mats_clone = mats.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_clone = vecs.detach().clone().requires_grad_(True)
+
     def test_inv_quad_log_det_many_vectors(self):
         # Forward pass
         actual_inv_quad = (
-            torch.cat([self.mats_var_clone[0].inverse().unsqueeze(0), self.mats_var_clone[1].inverse().unsqueeze(0)])
-            .matmul(self.vecs_var_clone)
-            .mul(self.vecs_var_clone)
+            torch.cat([mat.inverse().unsqueeze(0) for mat in self.mats_clone])
+            .matmul(self.vecs_clone)
+            .mul(self.vecs_clone)
             .sum(2)
             .sum(1)
         )
-        actual_log_det = torch.cat(
-            [self.mats_var_clone[0].det().log().unsqueeze(0), self.mats_var_clone[1].det().log().unsqueeze(0)]
-        )
-        with gpytorch.settings.num_trace_samples(1000):
-            nlv = NonLazyTensor(self.mats_var)
-            res_inv_quad, res_log_det = nlv.inv_quad_log_det(inv_quad_rhs=self.vecs_var, log_det=True)
-        self.assertTrue(approx_equal(res_inv_quad, actual_inv_quad, epsilon=1e-1))
-        self.assertTrue(approx_equal(res_log_det, actual_log_det, epsilon=1e-1))
+        actual_log_det = torch.cat([mat.logdet().unsqueeze(0) for mat in self.mats_clone])
+        with gpytorch.settings.num_trace_samples(2000):
+            non_lazy_tsr = NonLazyTensor(self.mats)
+            res_inv_quad, res_log_det = non_lazy_tsr.inv_quad_log_det(inv_quad_rhs=self.vecs, log_det=True)
+
+        self.assertEqual(res_inv_quad.shape, actual_inv_quad.shape)
+        self.assertEqual(res_log_det.shape, actual_log_det.shape)
+        self.assertLess(torch.max((res_inv_quad - actual_inv_quad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((res_log_det - actual_log_det).abs()).item(), 1e-1)
 
         # Backward
-        inv_quad_grad_output = torch.tensor([3, 4], dtype=torch.float)
-        log_det_grad_output = torch.tensor([4, 2], dtype=torch.float)
+        inv_quad_grad_output = torch.randn(5, dtype=torch.float)
+        log_det_grad_output = torch.randn(5, dtype=torch.float)
         actual_inv_quad.backward(gradient=inv_quad_grad_output)
         actual_log_det.backward(gradient=log_det_grad_output)
         res_inv_quad.backward(gradient=inv_quad_grad_output, retain_graph=True)
         res_log_det.backward(gradient=log_det_grad_output)
 
-        self.assertTrue(approx_equal(self.mats_var_clone.grad, self.mats_var.grad, epsilon=1e-1))
-        self.assertTrue(approx_equal(self.vecs_var_clone.grad, self.vecs_var.grad))
+        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
 
     def test_inv_quad_only_many_vectors(self):
         # Forward pass
-        res = NonLazyTensor(self.mats_var).inv_quad(self.vecs_var).sum()
+        res = NonLazyTensor(self.mats).inv_quad(self.vecs)
         actual = (
-            torch.cat([self.mats_var_clone[0].inverse().unsqueeze(0), self.mats_var_clone[1].inverse().unsqueeze(0)])
-            .matmul(self.vecs_var_clone)
-            .mul(self.vecs_var_clone)
+            torch.cat([mat.inverse().unsqueeze(0) for mat in self.mats_clone])
+            .matmul(self.vecs_clone)
+            .mul(self.vecs_clone)
             .sum(2)
             .sum(1)
-        ).sum()
-        self.assertTrue(approx_equal(res, actual, epsilon=1e-1))
+        )
+        self.assertEqual(res.shape, actual.shape)
+        self.assertLess(torch.max((res - actual).abs()).item(), 1e-1)
+
         # Backward
-        actual.backward()
-        res.backward()
+        actual.sum().backward()
+        res.sum().backward()
 
-        self.assertTrue(approx_equal(self.mats_var_clone.grad, self.mats_var.grad, epsilon=1e-1))
-        self.assertTrue(approx_equal(self.vecs_var_clone.grad, self.vecs_var.grad))
+        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-3)
+        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-3)
 
     def test_log_det_only(self):
         # Forward pass
-        with gpytorch.settings.num_trace_samples(1000):
-            res = NonLazyTensor(self.mats_var).log_det()
-        actual = torch.cat(
-            [self.mats_var_clone[0].det().log().unsqueeze(0), self.mats_var_clone[1].det().log().unsqueeze(0)]
+        with gpytorch.settings.num_trace_samples(2000):
+            res = NonLazyTensor(self.mats).log_det()
+        actual = torch.cat([mat.logdet().unsqueeze(0) for mat in self.mats_clone])
+        self.assertEqual(res.shape, actual.shape)
+        self.assertLess(torch.max((res - actual).abs()).item(), 1e-1)
+
+        # Backward
+        grad_output = torch.randn(5)
+        actual.backward(gradient=grad_output)
+        res.backward(gradient=grad_output)
+        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
+
+
+class TestInvQuadLogDetMultiBatch(unittest.TestCase):
+    def tearDown(self):
+        if hasattr(self, "rng_state"):
+            torch.set_rng_state(self.rng_state)
+
+    def setUp(self):
+        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
+            self.rng_state = torch.get_rng_state()
+            torch.manual_seed(0)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(0)
+            random.seed(0)
+
+        mats = torch.randn(2, 3, 4, 4)
+        mats = mats @ mats.transpose(-1, -2)
+        mats.div_(5).add_(torch.eye(4).view(1, 1, 4, 4))
+        vecs = torch.randn(2, 3, 4, 6)
+        self.mats = mats.detach().clone().requires_grad_(True)
+        self.mats_clone = mats.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_clone = vecs.detach().clone().requires_grad_(True)
+
+    def test_inv_quad_log_det_many_vectors(self):
+        # Forward pass
+        flattened_mats = self.mats_clone.view(-1, *self.mats_clone.shape[-2:])
+        actual_inv_quad = (
+            torch.cat([mat.inverse().unsqueeze(0) for mat in flattened_mats])
+            .view(self.mats_clone.shape)
+            .matmul(self.vecs_clone)
+            .mul(self.vecs_clone)
+            .sum(-2)
+            .sum(-1)
         )
-        self.assertTrue(approx_equal(res, actual, epsilon=1e-1))
+        actual_log_det = torch.cat([mat.logdet().unsqueeze(0) for mat in flattened_mats])
+        actual_log_det = actual_log_det.view(self.mats_clone.shape[:-2])
+
+        with gpytorch.settings.num_trace_samples(2000):
+            non_lazy_tsr = NonLazyTensor(self.mats)
+            res_inv_quad, res_log_det = non_lazy_tsr.inv_quad_log_det(inv_quad_rhs=self.vecs, log_det=True)
+
+        self.assertEqual(res_inv_quad.shape, actual_inv_quad.shape)
+        self.assertEqual(res_log_det.shape, actual_log_det.shape)
+        self.assertLess(torch.max((res_inv_quad - actual_inv_quad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((res_log_det - actual_log_det).abs()).item(), 1e-1)
+
+        # Backward
+        inv_quad_grad_output = torch.randn(2, 3, dtype=torch.float)
+        log_det_grad_output = torch.randn(2, 3, dtype=torch.float)
+        actual_inv_quad.backward(gradient=inv_quad_grad_output)
+        actual_log_det.backward(gradient=log_det_grad_output)
+        res_inv_quad.backward(gradient=inv_quad_grad_output, retain_graph=True)
+        res_log_det.backward(gradient=log_det_grad_output)
+
+        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
+
+    def test_inv_quad_only_many_vectors(self):
+        # Forward pass
+        res = NonLazyTensor(self.mats).inv_quad(self.vecs)
+        flattened_mats = self.mats_clone.view(-1, *self.mats_clone.shape[-2:])
+        actual = (
+            torch.cat([mat.inverse().unsqueeze(0) for mat in flattened_mats])
+            .view(self.mats_clone.shape)
+            .matmul(self.vecs_clone)
+            .mul(self.vecs_clone)
+            .sum(-2)
+            .sum(-1)
+        )
+        self.assertEqual(res.shape, actual.shape)
+        self.assertLess(torch.max((res - actual).abs()).item(), 1e-1)
+
+        # Backward
+        actual.sum().backward()
+        res.sum().backward()
+
+        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-3)
+        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-3)
+
+    def test_log_det_only(self):
+        # Forward pass
+        with gpytorch.settings.num_trace_samples(2000):
+            res = NonLazyTensor(self.mats).log_det()
+        flattened_mats = self.mats_clone.view(-1, *self.mats_clone.shape[-2:])
+        actual = torch.cat([mat.logdet().unsqueeze(0) for mat in flattened_mats])
+        actual = actual.view(self.mats_clone.shape[:-2])
+        self.assertEqual(res.shape, actual.shape)
+        self.assertLess(torch.max((res - actual).abs()).item(), 1e-1)
 
         # Backward
-        grad_output = torch.tensor([3, 4], dtype=torch.float)
+        grad_output = torch.randn(2, 3)
         actual.backward(gradient=grad_output)
         res.backward(gradient=grad_output)
-        self.assertTrue(approx_equal(self.mats_var_clone.grad, self.mats_var.grad, epsilon=1e-1))
+        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
 
 
 if __name__ == "__main__":
diff --git a/test/functions/test_matmul.py b/test/functions/test_matmul.py
index 9df14895..591ffd19 100644
--- a/test/functions/test_matmul.py
+++ b/test/functions/test_matmul.py
@@ -11,66 +11,87 @@ from test._utils import approx_equal
 
 class TestMatmulNonBatch(unittest.TestCase):
     def setUp(self):
-        mat = [[3, -1, 0], [-1, 3, 0], [0, 0, 3]]
-        vec = torch.randn(3, requires_grad=True)
-        vecs = torch.randn(3, 4, requires_grad=True)
-
-        self.mat_var = torch.tensor(mat, dtype=torch.float, requires_grad=True)
-        self.mat_var_copy = self.mat_var.clone().detach().requires_grad_(True)
-        self.vec_var = vec
-        self.vec_var_copy = vec.clone().detach().requires_grad_(True)
-        self.vecs_var = vecs
-        self.vecs_var_copy = vecs.clone().detach().requires_grad_(True)
+        mat = torch.tensor([[3, -1, 0], [-1, 3, 0], [0, 0, 3]], dtype=torch.float)
+        vec = torch.randn(3)
+        vecs = torch.randn(3, 4)
+        self.mat = mat.detach().clone().requires_grad_(True)
+        self.mat_copy = mat.detach().clone().requires_grad_(True)
+        self.vec = vec.detach().clone().requires_grad_(True)
+        self.vec_copy = vec.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_copy = vecs.detach().clone().requires_grad_(True)
 
     def test_matmul_vec(self):
         # Forward
-        res = NonLazyTensor(self.mat_var).matmul(self.vec_var)
-        actual = self.mat_var_copy.matmul(self.vec_var_copy)
+        res = NonLazyTensor(self.mat).matmul(self.vec)
+        actual = self.mat_copy.matmul(self.vec_copy)
         self.assertTrue(approx_equal(res, actual))
 
         # Backward
         grad_output = torch.randn(3)
         res.backward(gradient=grad_output)
         actual.backward(gradient=grad_output)
-        self.assertTrue(approx_equal(self.mat_var_copy.grad, self.mat_var.grad))
-        self.assertTrue(approx_equal(self.vec_var_copy.grad, self.vec_var.grad))
+        self.assertTrue(approx_equal(self.mat_copy.grad, self.mat.grad))
+        self.assertTrue(approx_equal(self.vec_copy.grad, self.vec.grad))
 
     def test_matmul_multiple_vecs(self):
         # Forward
-        res = NonLazyTensor(self.mat_var).matmul(self.vecs_var)
-        actual = self.mat_var_copy.matmul(self.vecs_var_copy)
+        res = NonLazyTensor(self.mat).matmul(self.vecs)
+        actual = self.mat_copy.matmul(self.vecs_copy)
         self.assertTrue(approx_equal(res, actual))
 
         # Backward
         grad_output = torch.randn(3, 4)
         res.backward(gradient=grad_output)
         actual.backward(gradient=grad_output)
-        self.assertTrue(approx_equal(self.mat_var_copy.grad, self.mat_var.grad))
-        self.assertTrue(approx_equal(self.vecs_var_copy.grad, self.vecs_var.grad))
+        self.assertTrue(approx_equal(self.mat_copy.grad, self.mat.grad))
+        self.assertTrue(approx_equal(self.vecs_copy.grad, self.vecs.grad))
 
 
 class TestMatmulBatch(unittest.TestCase):
     def setUp(self):
-        mats = [[[3, -1, 0], [-1, 3, 0], [0, 0, 3]], [[10, -2, 1], [-2, 10, 0], [1, 0, 10]]]
-        vecs = torch.randn(2, 3, 4, requires_grad=True)
+        mats = torch.randn(2, 5, 3)
+        vecs = torch.randn(2, 3, 4)
+        self.mats = mats.detach().clone().requires_grad_(True)
+        self.mats_copy = mats.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_copy = vecs.detach().clone().requires_grad_(True)
+
+    def test_matmul_multiple_vecs(self):
+        # Forward
+        res = NonLazyTensor(self.mats).matmul(self.vecs)
+        actual = self.mats_copy.matmul(self.vecs_copy)
+        self.assertTrue(approx_equal(res, actual))
 
-        self.mats_var = torch.tensor(mats, dtype=torch.float, requires_grad=True)
-        self.mats_var_clone = self.mats_var.clone().detach().requires_grad_(True)
-        self.vecs_var = vecs
-        self.vecs_var_copy = vecs.clone().detach().requires_grad_(True)
+        # Backward
+        grad_output = torch.randn(2, 5, 4)
+        res.backward(gradient=grad_output)
+        actual.backward(gradient=grad_output)
+        self.assertTrue(approx_equal(self.mats_copy.grad, self.mats.grad))
+        self.assertTrue(approx_equal(self.vecs_copy.grad, self.vecs.grad))
+
+
+class TestMatmulMultiBatch(unittest.TestCase):
+    def setUp(self):
+        mats = torch.randn(3, 4, 5, 6)
+        vecs = torch.randn(3, 4, 6, 2)
+        self.mats = mats.detach().clone().requires_grad_(True)
+        self.mats_copy = mats.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_copy = vecs.detach().clone().requires_grad_(True)
 
     def test_matmul_multiple_vecs(self):
         # Forward
-        res = NonLazyTensor(self.mats_var).matmul(self.vecs_var)
-        actual = self.mats_var_clone.matmul(self.vecs_var_copy)
+        res = NonLazyTensor(self.mats).matmul(self.vecs)
+        actual = self.mats_copy.matmul(self.vecs_copy)
         self.assertTrue(approx_equal(res, actual))
 
         # Backward
-        grad_output = torch.randn(2, 3, 4)
+        grad_output = torch.randn(3, 4, 5, 2)
         res.backward(gradient=grad_output)
         actual.backward(gradient=grad_output)
-        self.assertTrue(approx_equal(self.mats_var_clone.grad, self.mats_var.grad))
-        self.assertTrue(approx_equal(self.vecs_var_copy.grad, self.vecs_var.grad))
+        self.assertTrue(approx_equal(self.mats_copy.grad, self.mats.grad))
+        self.assertTrue(approx_equal(self.vecs_copy.grad, self.vecs.grad))
 
 
 if __name__ == "__main__":
diff --git a/test/functions/test_root_decomposition.py b/test/functions/test_root_decomposition.py
index 3f62c23e..b3ba59c3 100644
--- a/test/functions/test_root_decomposition.py
+++ b/test/functions/test_root_decomposition.py
@@ -5,52 +5,146 @@ from __future__ import unicode_literals
 
 import os
 import torch
+import random
 import unittest
 from gpytorch.lazy import NonLazyTensor
 from test._utils import approx_equal
 
 
 class TestRootDecomposition(unittest.TestCase):
+    def tearDown(self):
+        if hasattr(self, "rng_state"):
+            torch.set_rng_state(self.rng_state)
+
     def setUp(self):
-        if os.getenv("UNLOCK_SEED") is None or os.getenv("UNLOCK_SEED").lower() == "false":
+        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
             self.rng_state = torch.get_rng_state()
             torch.manual_seed(0)
-        mat = [
-            [5.0212, 0.5504, -0.1810, 1.5414, 2.9611],
-            [0.5504, 2.8000, 1.9944, 0.6208, -0.8902],
-            [-0.1810, 1.9944, 3.0505, 1.0790, -1.1774],
-            [1.5414, 0.6208, 1.0790, 2.9430, 0.4170],
-            [2.9611, -0.8902, -1.1774, 0.4170, 3.3208],
-        ]
-        self.mat_var = torch.tensor(mat, requires_grad=True)
-        self.mat_var_clone = self.mat_var.clone().detach().requires_grad_(True)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(0)
+            random.seed(0)
 
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
+        mat = torch.randn(4, 4)
+        mat = mat @ mat.transpose(-1, -2)
+        mat.div_(5).add_(torch.eye(4))
+        self.mat = mat.detach().clone().requires_grad_(True)
+        self.mat_clone = mat.detach().clone().requires_grad_(True)
 
     def test_root_decomposition(self):
         # Forward
-        root = NonLazyTensor(self.mat_var).root_decomposition()
+        root = NonLazyTensor(self.mat).root_decomposition()
         res = root.matmul(root.transpose(-1, -2))
-        self.assertTrue(approx_equal(res, self.mat_var))
+        self.assertTrue(approx_equal(res, self.mat))
 
         # Backward
         res.trace().backward()
-        self.mat_var_clone.trace().backward()
-        self.assertTrue(approx_equal(self.mat_var.grad, self.mat_var_clone.grad))
+        self.mat_clone.trace().backward()
+        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
 
     def test_root_inv_decomposition(self):
         # Forward
-        root = NonLazyTensor(self.mat_var).root_inv_decomposition()
+        probe_vectors = torch.randn(4, 5)
+        test_vectors = torch.randn(4, 5)
+        root = NonLazyTensor(self.mat).root_inv_decomposition(probe_vectors, test_vectors)
         res = root.matmul(root.transpose(-1, -2))
-        actual = self.mat_var_clone.inverse()
+        actual = self.mat_clone.inverse()
         self.assertTrue(approx_equal(res, actual))
 
         # Backward
         res.trace().backward()
         actual.trace().backward()
-        self.assertTrue(approx_equal(self.mat_var.grad, self.mat_var_clone.grad))
+        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
+
+
+class TestRootDecompositionBatch(unittest.TestCase):
+    def tearDown(self):
+        if hasattr(self, "rng_state"):
+            torch.set_rng_state(self.rng_state)
+
+    def setUp(self):
+        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
+            self.rng_state = torch.get_rng_state()
+            torch.manual_seed(0)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(0)
+            random.seed(0)
+
+        mat = torch.randn(3, 4, 4)
+        mat = mat @ mat.transpose(-1, -2)
+        mat.div_(5).add_(torch.eye(4).unsqueeze_(0))
+        self.mat = mat.detach().clone().requires_grad_(True)
+        self.mat_clone = mat.detach().clone().requires_grad_(True)
+
+    def test_root_decomposition(self):
+        # Forward
+        root = NonLazyTensor(self.mat).root_decomposition()
+        res = root.matmul(root.transpose(-1, -2))
+        self.assertTrue(approx_equal(res, self.mat))
+
+        # Backward
+        sum([mat.trace() for mat in res]).backward()
+        sum([mat.trace() for mat in self.mat_clone]).backward()
+        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
+
+    def test_root_inv_decomposition(self):
+        # Forward
+        probe_vectors = torch.randn(3, 4, 5)
+        test_vectors = torch.randn(3, 4, 5)
+        root = NonLazyTensor(self.mat).root_inv_decomposition(probe_vectors, test_vectors)
+        res = root.matmul(root.transpose(-1, -2))
+        actual = torch.cat([mat.inverse().unsqueeze(0) for mat in self.mat_clone])
+        self.assertTrue(approx_equal(res, actual))
+
+        # Backward
+        sum([mat.trace() for mat in res]).backward()
+        sum([mat.trace() for mat in actual]).backward()
+        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
+
+
+class TestRootDecompositionMultiBatch(unittest.TestCase):
+    def tearDown(self):
+        if hasattr(self, "rng_state"):
+            torch.set_rng_state(self.rng_state)
+
+    def setUp(self):
+        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
+            self.rng_state = torch.get_rng_state()
+            torch.manual_seed(0)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(0)
+            random.seed(0)
+
+        mat = torch.randn(2, 3, 4, 4)
+        mat = mat @ mat.transpose(-1, -2)
+        mat.div_(5).add_(torch.eye(4).view(1, 1, 4, 4))
+        self.mat = mat.detach().clone().requires_grad_(True)
+        self.mat_clone = mat.detach().clone().requires_grad_(True)
+
+    def test_root_decomposition(self):
+        # Forward
+        root = NonLazyTensor(self.mat).root_decomposition()
+        res = root.matmul(root.transpose(-1, -2))
+        self.assertTrue(approx_equal(res, self.mat))
+
+        # Backward
+        sum([mat.trace() for mat in res.view(-1, *self.mat.shape[-2:])]).backward()
+        sum([mat.trace() for mat in self.mat_clone.view(-1, *self.mat.shape[-2:])]).backward()
+        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
+
+    def test_root_inv_decomposition(self):
+        # Forward
+        probe_vectors = torch.randn(2, 3, 4, 5)
+        test_vectors = torch.randn(2, 3, 4, 5)
+        root = NonLazyTensor(self.mat).root_inv_decomposition(probe_vectors, test_vectors)
+        res = root.matmul(root.transpose(-1, -2))
+        flattened_mats = self.mat_clone.view(-1, *self.mat_clone.shape[-2:])
+        actual = torch.cat([mat.inverse().unsqueeze(0) for mat in flattened_mats]).view_as(self.mat_clone)
+        self.assertTrue(approx_equal(res, actual))
+
+        # Backward
+        sum([mat.trace() for mat in res.view(-1, *self.mat.shape[-2:])]).backward()
+        sum([mat.trace() for mat in actual.view(-1, *self.mat.shape[-2:])]).backward()
+        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
 
 
 if __name__ == "__main__":
diff --git a/test/lazy/_lazy_tensor_test_case.py b/test/lazy/_lazy_tensor_test_case.py
index f51af5bb..a50c1f01 100644
--- a/test/lazy/_lazy_tensor_test_case.py
+++ b/test/lazy/_lazy_tensor_test_case.py
@@ -8,6 +8,7 @@ import torch
 import os
 import random
 from abc import abstractmethod
+from itertools import product
 from test._utils import approx_equal
 
 
@@ -214,7 +215,7 @@ class RectangularBatchLazyTensorTestCase(object):
         lazy_tensor_copy = lazy_tensor.clone()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
 
-        test_vector = torch.randn(lazy_tensor.size(0), lazy_tensor.size(-1), 5)
+        test_vector = torch.randn(*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)
@@ -246,44 +247,37 @@ class RectangularBatchLazyTensorTestCase(object):
         res = lazy_tensor[:, 0:2].evaluate()
         actual = evaluated[:, 0:2]
         self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
-        res = lazy_tensor[1, :1, :2].evaluate()
-        actual = evaluated[1, :1, :2]
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
-        res = lazy_tensor[1, 1, :2]
-        actual = evaluated[1, 1, :2]
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
-        res = lazy_tensor[1, :1, 2]
-        actual = evaluated[1, :1, 2]
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+
+        for batch_index in product([1, slice(0, 2, None)], repeat=(lazy_tensor.dim() - 2)):
+            res = lazy_tensor.__getitem__((*batch_index, slice(0, 1, None), slice(0, 2, None))).evaluate()
+            actual = evaluated.__getitem__((*batch_index, slice(0, 1, None), slice(0, 2, None)))
+            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            res = lazy_tensor.__getitem__((*batch_index, 1, slice(0, 2, None)))
+            actual = evaluated.__getitem__((*batch_index, 1, slice(0, 2, None)))
+            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            res = lazy_tensor.__getitem__((*batch_index, slice(1, None, None), 2))
+            actual = evaluated.__getitem__((*batch_index, slice(1, None, None), 2))
+            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
 
     def test_getitem_tensor_index(self):
         lazy_tensor = self.create_lazy_tensor()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
 
-        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 1]))
-        res, actual = lazy_tensor[index], evaluated[index]
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
-        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
-        res, actual = lazy_tensor[index], evaluated[index]
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
-        index = (torch.tensor([0, 1, 1]), slice(None, None, None), torch.tensor([0, 1, 2]))
-        res, actual = lazy_tensor[index], evaluated[index]
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
-        index = (slice(None, None, None), torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]))
-        res, actual = lazy_tensor[index], evaluated[index]
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
-        index = (torch.tensor([0, 0, 1, 1]), slice(None, None, None), slice(None, None, None))
-        res, actual = lazy_tensor[index].evaluate(), evaluated[index]
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
-        index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]), torch.tensor([0, 0, 1, 1]))
-        res, actual = lazy_tensor[index], evaluated[index]
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
-        index = (torch.tensor([0, 1, 1, 0]), torch.tensor([0, 1, 0, 2]), slice(None, None, None))
-        res, actual = lazy_tensor[index], evaluated[index]
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
-        index = (torch.tensor([0, 0, 1, 0]), slice(None, None, None), torch.tensor([0, 0, 1, 1]))
-        res, actual = lazy_tensor[index], evaluated[index]
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+        for batch_index in product([
+            torch.tensor([0, 1, 1, 0]), slice(None, None, None)
+        ], repeat=(lazy_tensor.dim() - 2)):
+            index = (*batch_index, torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 1]))
+            res, actual = lazy_tensor[index], evaluated[index]
+            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            index = (*batch_index, torch.tensor([0, 1, 0, 2]), slice(None, None, None))
+            res, actual = lazy_tensor[index], evaluated[index]
+            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            index = (*batch_index, slice(None, None, None), torch.tensor([0, 1, 2, 1]))
+            res, actual = lazy_tensor[index], evaluated[index]
+            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            index = (*batch_index, slice(None, None, None), slice(None, None, None))
+            res, actual = lazy_tensor[index].evaluate(), evaluated[index]
+            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
 
 
 class BatchLazyTensorTestCase(RectangularBatchLazyTensorTestCase):
@@ -317,47 +311,48 @@ class BatchLazyTensorTestCase(RectangularBatchLazyTensorTestCase):
 
         other_diag = torch.tensor(1.5)
         res = lazy_tensor.add_diag(other_diag).evaluate()
-        actual = evaluated + torch.eye(evaluated.size(-1)).unsqueeze(0).repeat(evaluated.size(0), 1, 1).mul(1.5)
+        actual = evaluated + torch.eye(evaluated.size(-1)).view(
+            *[1 for _ in range(lazy_tensor.dim() - 2)], evaluated.size(-1), evaluated.size(-1)
+        ).repeat(*lazy_tensor.batch_shape, 1, 1).mul(1.5)
         self.assertTrue(approx_equal(res, actual))
 
         other_diag = torch.tensor([1.5])
         res = lazy_tensor.add_diag(other_diag).evaluate()
-        actual = evaluated + torch.eye(evaluated.size(-1)).unsqueeze(0).repeat(evaluated.size(0), 1, 1).mul(1.5)
+        actual = evaluated + torch.eye(evaluated.size(-1)).view(
+            *[1 for _ in range(lazy_tensor.dim() - 2)], evaluated.size(-1), evaluated.size(-1)
+        ).repeat(*lazy_tensor.batch_shape, 1, 1).mul(1.5)
         self.assertTrue(approx_equal(res, actual))
 
         other_diag = torch.randn(lazy_tensor.size(-1)).pow(2)
         res = lazy_tensor.add_diag(other_diag).evaluate()
-        actual = evaluated + other_diag.diag().unsqueeze(0).repeat(evaluated.size(0), 1, 1)
-        self.assertTrue(approx_equal(res, actual))
-
-        other_diag = torch.randn(1, lazy_tensor.size(-1)).pow(2)
-        res = lazy_tensor.add_diag(other_diag).evaluate()
-        actual = evaluated + other_diag[0].diag().unsqueeze(0).repeat(evaluated.size(0), 1, 1)
-        self.assertTrue(approx_equal(res, actual))
-
-        other_diag = torch.randn(lazy_tensor.size(0), lazy_tensor.size(-1)).pow(2)
-        res = lazy_tensor.add_diag(other_diag).evaluate()
-        actual = evaluated + torch.cat([other_diag[i].diag().unsqueeze(0) for i in range(lazy_tensor.size(0))])
+        actual = evaluated + other_diag.diag().repeat(*lazy_tensor.batch_shape, 1, 1)
         self.assertTrue(approx_equal(res, actual))
 
-        other_diag = torch.randn(lazy_tensor.size(0), 1).pow(2)
-        res = lazy_tensor.add_diag(other_diag).evaluate()
-        actual = evaluated + torch.cat(
-            [torch.eye(lazy_tensor.size(-1)).mul(other_diag[i, 0]).unsqueeze(0) for i in range(lazy_tensor.size(0))]
-        )
-        self.assertTrue(approx_equal(res, actual))
+        for sizes in product([1, None], repeat=(lazy_tensor.dim() - 2)):
+            batch_shape = [
+                lazy_tensor.batch_shape[i] if size is None else size for i, size in enumerate(sizes)
+            ]
+            other_diag = torch.randn(*batch_shape, lazy_tensor.size(-1)).pow(2)
+            res = lazy_tensor.add_diag(other_diag).evaluate()
+            actual = evaluated.clone().detach()
+            for i in range(other_diag.size(-1)):
+                actual[..., i, i] = actual[..., i, i] + other_diag[..., i]
+            self.assertTrue(approx_equal(res, actual))
 
     def test_inv_matmul_matrix(self):
         lazy_tensor = self.create_lazy_tensor()
         lazy_tensor_copy = lazy_tensor.clone()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+        flattened_evaluated = evaluated.view(-1, *lazy_tensor.matrix_shape)
 
-        test_vector = torch.randn(lazy_tensor.size(0), lazy_tensor.size(-1), 5)
+        test_vector = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5)
         with gpytorch.settings.max_cg_iterations(200):
             res = lazy_tensor.inv_matmul(test_vector)
-        actual = torch.cat(
-            [evaluated[i].inverse().matmul(test_vector[i]).unsqueeze(0) for i in range(lazy_tensor.size(0))]
-        )
+        actual = torch.cat([
+            flattened_evaluated[i].inverse().matmul(test_vector[i]).unsqueeze(0)
+            for i in range(lazy_tensor.batch_shape.numel())
+        ])
+        actual = actual.view_as(test_vector)
         self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
 
         grad = torch.randn_like(res)
@@ -372,9 +367,11 @@ class BatchLazyTensorTestCase(RectangularBatchLazyTensorTestCase):
     def test_diag(self):
         lazy_tensor = self.create_lazy_tensor()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
+        flattened_evaluated = evaluated.view(-1, *lazy_tensor.matrix_shape)
 
         res = lazy_tensor.diag()
-        actual = torch.stack([evaluated[i].diag() for i in range(evaluated.size(0))])
+        actual = torch.stack([flattened_evaluated[i].diag() for i in range(flattened_evaluated.size(0))])
+        actual = actual.view(*lazy_tensor.batch_shape, -1)
         self.assertEqual(res.size(), lazy_tensor.size()[:-1])
         self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
 
@@ -382,8 +379,9 @@ class BatchLazyTensorTestCase(RectangularBatchLazyTensorTestCase):
         # Forward
         lazy_tensor = self.create_lazy_tensor()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
+        flattened_evaluated = evaluated.view(-1, *lazy_tensor.matrix_shape)
 
-        vecs = torch.randn(lazy_tensor.size(0), lazy_tensor.size(1), 3, requires_grad=True)
+        vecs = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 3, requires_grad=True)
         vecs_copy = vecs.clone()
 
         with gpytorch.settings.num_trace_samples(128):
@@ -392,15 +390,19 @@ class BatchLazyTensorTestCase(RectangularBatchLazyTensorTestCase):
 
         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))
+                flattened_evaluated[i].inverse().matmul(vecs_copy[i]).mul(vecs_copy[i]).sum().unsqueeze(0)
+                for i in range(lazy_tensor.batch_shape.numel())
             ]
         )
         actual = actual_inv_quad + torch.cat(
-            [torch.logdet(evaluated[i]).unsqueeze(0) for i in range(lazy_tensor.size(0))]
+            [torch.logdet(flattened_evaluated[i]).unsqueeze(0) for i in range(lazy_tensor.batch_shape.numel())]
         )
-        diffs = (res - actual).abs() / actual.abs().clamp(1, 1e10)
-        for i in range(lazy_tensor.size(0)):
+        actual = actual.view(*lazy_tensor.batch_shape)
+
+        self.assertEqual(actual.shape, res.shape)
+
+        diffs = (res - actual).abs() / actual.abs().clamp(1, 1e10).view(-1)
+        for i in range(diffs.numel()):
             self.assertLess(diffs[i].item(), 15e-2)
 
     def test_sample(self):
diff --git a/test/utils/test_lanczos.py b/test/utils/test_lanczos.py
index ecc371c9..fb930fd6 100644
--- a/test/utils/test_lanczos.py
+++ b/test/utils/test_lanczos.py
@@ -17,7 +17,7 @@ class TestLanczos(unittest.TestCase):
         matrix.div_(matrix.norm())
         matrix.add_(torch.ones(matrix.size(-1)).mul(1e-6).diag())
         q_mat, t_mat = lanczos_tridiag(
-            matrix.matmul, max_iter=size, dtype=matrix.dtype, device=matrix.device, n_dims=matrix.size(-1)
+            matrix.matmul, max_iter=size, dtype=matrix.dtype, device=matrix.device, matrix_shape=matrix.shape
         )
 
         approx = q_mat.matmul(t_mat).matmul(q_mat.transpose(-1, -2))
diff --git a/test/utils/test_pivoted_cholesky.py b/test/utils/test_pivoted_cholesky.py
index 814e2294..b3b29b0a 100644
--- a/test/utils/test_pivoted_cholesky.py
+++ b/test/utils/test_pivoted_cholesky.py
@@ -97,5 +97,69 @@ class TestPivotedCholeskyBatch(unittest.TestCase):
         self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))
 
 
+class TestPivotedCholeskyMultiBatch(unittest.TestCase):
+    def test_pivoted_cholesky(self):
+        size = 100
+        train_x = torch.cat([
+            torch.linspace(0, 1, size).unsqueeze(0),
+            torch.linspace(0, 0.5, size).unsqueeze(0),
+            torch.linspace(0, 0.25, size).unsqueeze(0),
+            torch.linspace(0, 1.25, size).unsqueeze(0),
+            torch.linspace(0, 1.5, size).unsqueeze(0),
+            torch.linspace(0, 1, size).unsqueeze(0),
+            torch.linspace(0, 0.5, size).unsqueeze(0),
+            torch.linspace(0, 0.25, size).unsqueeze(0),
+            torch.linspace(0, 1.25, size).unsqueeze(0),
+            torch.linspace(0, 1.25, size).unsqueeze(0),
+            torch.linspace(0, 1.5, size).unsqueeze(0),
+            torch.linspace(0, 1, size).unsqueeze(0),
+        ], 0).unsqueeze(-1)
+        covar_matrix = RBFKernel()(train_x, train_x).evaluate().view(2, 2, 3, size, size)
+        piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
+        covar_approx = piv_chol.transpose(-1, -2).matmul(piv_chol)
+
+        self.assertTrue(approx_equal(covar_approx, covar_matrix, 2e-4))
+
+    def test_solve(self):
+        size = 100
+        train_x = torch.cat([
+            torch.linspace(0, 1, size).unsqueeze(0),
+            torch.linspace(0, 0.5, size).unsqueeze(0),
+            torch.linspace(0, 0.25, size).unsqueeze(0),
+            torch.linspace(0, 1.25, size).unsqueeze(0),
+            torch.linspace(0, 1.5, size).unsqueeze(0),
+            torch.linspace(0, 1, size).unsqueeze(0),
+            torch.linspace(0, 0.5, size).unsqueeze(0),
+            torch.linspace(0, 0.25, size).unsqueeze(0),
+            torch.linspace(0, 1.25, size).unsqueeze(0),
+            torch.linspace(0, 1.25, size).unsqueeze(0),
+            torch.linspace(0, 1.5, size).unsqueeze(0),
+            torch.linspace(0, 1, size).unsqueeze(0),
+        ], 0).unsqueeze(-1)
+        covar_matrix = RBFKernel()(train_x, train_x).evaluate().view(2, 2, 3, size, size)
+        piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
+        woodbury_factor = pivoted_cholesky.woodbury_factor(piv_chol, torch.ones(2, 2, 3, 100))
+
+        rhs_vector = torch.randn(2, 2, 3, 100, 5)
+        shifted_covar_matrix = covar_matrix + torch.eye(size)
+        real_solve = torch.cat([
+            shifted_covar_matrix[0, 0, 0].inverse().matmul(rhs_vector[0, 0, 0]).unsqueeze(0),
+            shifted_covar_matrix[0, 0, 1].inverse().matmul(rhs_vector[0, 0, 1]).unsqueeze(0),
+            shifted_covar_matrix[0, 0, 2].inverse().matmul(rhs_vector[0, 0, 2]).unsqueeze(0),
+            shifted_covar_matrix[0, 1, 0].inverse().matmul(rhs_vector[0, 1, 0]).unsqueeze(0),
+            shifted_covar_matrix[0, 1, 1].inverse().matmul(rhs_vector[0, 1, 1]).unsqueeze(0),
+            shifted_covar_matrix[0, 1, 2].inverse().matmul(rhs_vector[0, 1, 2]).unsqueeze(0),
+            shifted_covar_matrix[1, 0, 0].inverse().matmul(rhs_vector[1, 0, 0]).unsqueeze(0),
+            shifted_covar_matrix[1, 0, 1].inverse().matmul(rhs_vector[1, 0, 1]).unsqueeze(0),
+            shifted_covar_matrix[1, 0, 2].inverse().matmul(rhs_vector[1, 0, 2]).unsqueeze(0),
+            shifted_covar_matrix[1, 1, 0].inverse().matmul(rhs_vector[1, 1, 0]).unsqueeze(0),
+            shifted_covar_matrix[1, 1, 1].inverse().matmul(rhs_vector[1, 1, 1]).unsqueeze(0),
+            shifted_covar_matrix[1, 1, 2].inverse().matmul(rhs_vector[1, 1, 2]).unsqueeze(0),
+        ], 0).view_as(rhs_vector)
+        approx_solve = pivoted_cholesky.woodbury_solve(rhs_vector, piv_chol, woodbury_factor, torch.ones(2, 3, 100))
+
+        self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))
+
+
 if __name__ == "__main__":
     unittest.main()