Repository: cornellius-gp/gpytorch
Issue: KronckerProductLazyVariable only works for square and symmetric matrices

Easy enough to reproduce:

```python
a = torch.randn(5, 3)
b = torch.randn(2, 2)
kplv = KroneckerProductLazyVariable(NonLazyVariable(a), NonLazyVariable(b))
res = kplv._matmul(torch.randn(3 * 2)) # fails
```

This is blocking MultitaskKernel from working properly, so I'm trying to solve this now

Changed files:
test/lazy/test_kronecker_product_lazy_variable.py, test/util/test_linear_cg.py, gpytorch/kernels/linear_kernel.py, gpytorch/kernels/index_kernel.py, gpytorch/likelihoods/softmax_likelihood.py, gpytorch/lazy/lazy_evaluated_kernel_variable.py, gpytorch/kernels/kernel.py, gpytorch/utils/linear_cg.py, gpytorch/lazy/kronecker_product_lazy_variable.py

diff --git a/gpytorch/kernels/index_kernel.py b/gpytorch/kernels/index_kernel.py
index 7da45509..c84d3b00 100644
--- a/gpytorch/kernels/index_kernel.py
+++ b/gpytorch/kernels/index_kernel.py
@@ -15,19 +15,10 @@ def _eval_covar_matrix(covar_factor, log_var):
 class IndexKernel(Kernel):
     def __init__(self, n_tasks, rank=1, prior=None, active_dims=None):
         if active_dims is not None and len(active_dims) > 1:
-            raise ValueError(
-                "Index must be with respect to a single column. Received {}".format(
-                    active_dims
-                )
-            )
+            raise ValueError("Index must be with respect to a single column. Received {}".format(active_dims))
         super(IndexKernel, self).__init__(active_dims=active_dims)
-        self.register_parameter(
-            name="covar_factor",
-            parameter=torch.nn.Parameter(torch.randn(n_tasks, rank)),
-        )
-        self.register_parameter(
-            name="log_var", parameter=torch.nn.Parameter(torch.randn(n_tasks))
-        )
+        self.register_parameter(name="covar_factor", parameter=torch.nn.Parameter(torch.randn(n_tasks, rank)))
+        self.register_parameter(name="log_var", parameter=torch.nn.Parameter(torch.randn(n_tasks)))
         if prior is not None:
             self.register_derived_prior(
                 name="IndexKernelPrior",
@@ -40,9 +31,5 @@ class IndexKernel(Kernel):
         covar_matrix = _eval_covar_matrix(self.covar_factor, self.log_var)
         if covar_matrix.ndimension() == 2:
             covar_matrix = covar_matrix.unsqueeze(0)
-        res = InterpolatedLazyVariable(
-            base_lazy_variable=covar_matrix,
-            left_interp_indices=i1,
-            right_interp_indices=i2,
-        )
+        res = InterpolatedLazyVariable(base_lazy_variable=covar_matrix, left_interp_indices=i1, right_interp_indices=i2)
         return res
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index 79218f08..e3a2a5cc 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -37,6 +37,7 @@ class Kernel(Module):
         >>> lazy_covar_matrix = covar_module(x1) # Returns a RootLazyVariable
         >>> tensor_covar_matrix = lazy_covar_matrix.evaluate() # Gets the actual tensor for this kernel matrix
     """
+
     def __init__(
         self,
         has_lengthscale=False,
@@ -122,6 +123,7 @@ class AdditiveKernel(Kernel):
         >>> x1 = torch.randn(50, 2)
         >>> additive_kernel_matrix = covar_module(x1)
     """
+
     def __init__(self, *kernels):
         super(AdditiveKernel, self).__init__()
         self.kernels = ModuleList(kernels)
@@ -143,6 +145,7 @@ class ProductKernel(Kernel):
         >>> x1 = torch.randn(50, 2)
         >>> kernel_matrix = covar_module(x1) # The RBF Kernel already decomposes multiplicatively, so this is foolish!
     """
+
     def __init__(self, *kernels):
         super(ProductKernel, self).__init__()
         self.kernels = ModuleList(kernels)
diff --git a/gpytorch/kernels/linear_kernel.py b/gpytorch/kernels/linear_kernel.py
index 8b32a413..74be52e4 100644
--- a/gpytorch/kernels/linear_kernel.py
+++ b/gpytorch/kernels/linear_kernel.py
@@ -19,6 +19,7 @@ class LinearKernel(Kernel):
     :math:`Kv` can be computed as :math:`Kv=X(X^{\\top}v)`, where the base multiply :math:`Xv` takes only :math:`O(nd)`
     time and space.
     """
+
     def __init__(
         self,
         num_dimensions,
diff --git a/gpytorch/lazy/kronecker_product_lazy_variable.py b/gpytorch/lazy/kronecker_product_lazy_variable.py
index ed68d217..84888e84 100644
--- a/gpytorch/lazy/kronecker_product_lazy_variable.py
+++ b/gpytorch/lazy/kronecker_product_lazy_variable.py
@@ -21,15 +21,45 @@ def _matmul(lazy_vars, rhs):
     for lazy_var in list(lazy_vars)[::-1]:
         if is_batch:
             n_batch = res.size(0)
-            res = res.view(n_batch, lazy_var.size(-1), -1)
-            factor = lazy_var._matmul(res)
-            n_batch = factor.size(0)
-            factor = factor.view(n_batch, lazy_var.size(-1), -1, n_cols).transpose(-3, -2)
+            res = (
+                res.transpose(-2, -1)
+                .contiguous()
+                .view(n_batch, n_cols, lazy_var.size(-1), -1)
+                .transpose(0, 1)
+                .contiguous()
+            )
+            factor = lazy_var._matmul(res).permute(1, 3, 2, 0)
             res = factor.contiguous().view(n_batch, -1, n_cols)
         else:
-            res = res.view(lazy_var.size(-1), -1)
+            res = res.t().contiguous().view(n_cols, lazy_var.size(-1), -1)
             factor = lazy_var._matmul(res)
-            factor = factor.contiguous().view(lazy_var.size(-1), -1, n_cols).transpose(-3, -2)
+            if factor.ndimension() == 3:
+                factor = factor.permute(2, 1, 0)
+            else:
+                factor = factor.transpose(-2, -1)
+            res = factor.contiguous().view(-1, n_cols)
+    return res
+
+
+def _t_matmul(lazy_vars, rhs):
+    res = rhs.contiguous()
+    is_batch = rhs.ndimension() == 3
+    n_cols = rhs.size(-1)
+    for lazy_var in list(lazy_vars)[::-1]:
+        if is_batch:
+            n_batch = res.size(0)
+            res = (
+                res.transpose(-2, -1)
+                .contiguous()
+                .view(n_batch, n_cols, lazy_var.size(-2), -1)
+                .transpose(0, 1)
+                .contiguous()
+            )
+            factor = lazy_var._t_matmul(res).permute(1, 3, 2, 0)
+            res = factor.contiguous().view(n_batch, -1, n_cols)
+        else:
+            res = res.t().contiguous().view(n_cols, lazy_var.size(-2), -1)
+            factor = lazy_var._t_matmul(res).permute(2, 1, 0)
             res = factor.contiguous().view(-1, n_cols)
     return res
 
@@ -52,6 +82,17 @@ class KroneckerProductLazyVariable(LazyVariable):
             res = res.squeeze(-1)
         return res
 
+    def _t_matmul(self, rhs):
+        is_vec = rhs.ndimension() == 1
+        if is_vec:
+            rhs = rhs.unsqueeze(-1)
+
+        res = _t_matmul(self.lazy_vars, rhs.contiguous())
+
+        if is_vec:
+            res = res.squeeze(-1)
+        return res
+
     def _quad_form_derivative(self, left_vecs, right_vecs):
         if left_vecs.ndimension() == 1:
             left_vecs = left_vecs.unsqueeze(1)
@@ -60,61 +101,81 @@ class KroneckerProductLazyVariable(LazyVariable):
         left_vecs = left_vecs.contiguous()
         right_vecs = right_vecs.contiguous()
         is_batch = left_vecs.ndimension() == 3
-        sizes = [lazy_var.size(-1) for lazy_var in self.lazy_vars]
+        row_sizes = [lazy_var.size(-2) for lazy_var in self.lazy_vars]
+        col_sizes = [lazy_var.size(-1) for lazy_var in self.lazy_vars]
 
         res = []
         if not is_batch:
             _, s = left_vecs.size()
 
+            # We need to call the derivative routine for each of the sub matrices
             for i, lazy_var in enumerate(self.lazy_vars):
-                m_left = _prod(sizes[i + 1 :])
-                m_right = _prod(sizes[:i])
-                m_i = sizes[i]
-
-                right_vecs_i = right_vecs.view(m_left, m_i, m_right, s).transpose(0, 1).contiguous()
-                right_vecs_i = right_vecs_i.view(m_i, m_left * m_right * s)
-
-                left_vecs_i = left_vecs.view(m_left, m_i * m_right, s)
+                m_left_prev = _prod(row_sizes[i + 1 :])
+                m_left_next = _prod(row_sizes[:i])
+                m_left_i = row_sizes[i]
+                m_right_prev = _prod(col_sizes[i + 1 :])
+                m_right_next = _prod(col_sizes[:i])
+                m_right_i = col_sizes[i]
+
+                # The right vectors are the same for each of the sub matrices
+                # Here we're reshaping them so that the batch dimension corresponds to the correct vectors
+                right_vecs_i = right_vecs.view(m_right_prev, m_right_i, m_right_next, s).transpose(0, 1).contiguous()
+                right_vecs_i = right_vecs_i.view(m_right_i, m_right_prev * m_right_next * s)
+
+                # The left vectors need to be multiplied by all other sub matrices in the kronecker product
+                # First multiply by all the matrices in the KP that are to follow this matrix
                 if i != len(self.lazy_vars) - 1:
-                    left_vecs_i = left_vecs_i.view(m_left, m_i * m_right * s)
-                    left_vecs_i = _matmul(self.lazy_vars[i + 1 :], left_vecs_i)
-                    left_vecs_i = left_vecs_i.view(m_left, m_i * m_right, s)
-
-                left_vecs_i = left_vecs_i.view(m_left, m_i, m_right, s).transpose(0, 1)
+                    left_vecs_i = left_vecs.view(m_left_prev, m_left_i * m_left_next * s)
+                    left_vecs_i = _t_matmul(self.lazy_vars[i + 1 :], left_vecs_i)
+                    m_left_prev = m_right_prev
+                    left_vecs_i = left_vecs_i.view(m_left_prev, m_left_i * m_left_next, s)
+                else:
+                    left_vecs_i = left_vecs.view(m_left_prev, m_left_i * m_left_next, s)
+
+                # Next multiply by all the previous matrices in the KP
+                left_vecs_i = left_vecs_i.view(m_left_prev, m_left_i, m_left_next, s).transpose(0, 1)
                 if i != 0:
-                    left_vecs_i = left_vecs_i.transpose(0, 2).contiguous().view(m_right, m_left * m_i * s)
-                    left_vecs_i = _matmul(self.lazy_vars[:i], left_vecs_i)
-                    left_vecs_i = left_vecs_i.view(m_right, m_left, m_i, s).transpose(0, 2)
+                    left_vecs_i = left_vecs_i.transpose(0, 2).contiguous().view(m_left_next, m_left_prev * m_left_i * s)
+                    left_vecs_i = _t_matmul(self.lazy_vars[:i], left_vecs_i)
+                    m_left_next = m_right_next
+                    left_vecs_i = left_vecs_i.view(m_left_next, m_left_prev, m_left_i, s).transpose(0, 2)
 
-                left_vecs_i = left_vecs_i.contiguous().view(m_i, m_left * m_right * s).contiguous()
+                # Now reshape the left vectors
+                left_vecs_i = left_vecs_i.contiguous().view(m_left_i, m_left_prev * m_left_next * s).contiguous()
                 res = res + list(lazy_var._quad_form_derivative(left_vecs_i, right_vecs_i))
         else:
             batch_size, _, s = left_vecs.size()
 
             for i, lazy_var in enumerate(self.lazy_vars):
-                m_left = _prod(sizes[i + 1 :])
-                m_right = _prod(sizes[:i])
-                m_i = sizes[i]
-
-                right_vecs_i = right_vecs.view(batch_size, m_left, m_i, m_right, s)
+                m_right_prev = _prod(col_sizes[i + 1 :])
+                m_right_next = _prod(col_sizes[:i])
+                m_right_i = col_sizes[i]
+                m_left_prev = _prod(row_sizes[i + 1 :])
+                m_left_next = _prod(row_sizes[:i])
+                m_left_i = row_sizes[i]
+
+                right_vecs_i = right_vecs.view(batch_size, m_right_prev, m_right_i, m_right_next, s)
                 right_vecs_i = right_vecs_i.transpose(1, 2).contiguous()
-                right_vecs_i = right_vecs_i.view(batch_size, m_i, m_left * m_right * s)
+                right_vecs_i = right_vecs_i.view(batch_size, m_right_i, m_right_prev * m_right_next * s)
 
-                left_vecs_i = left_vecs.view(batch_size, m_left, m_i * m_right, s)
                 if i != len(self.lazy_vars) - 1:
-                    left_vecs_i = left_vecs_i.view(batch_size, m_left, m_i * m_right * s)
-                    left_vecs_i = _matmul(self.lazy_vars[i + 1 :], left_vecs_i)
-                    left_vecs_i = left_vecs_i.contiguous().view(batch_size, m_left, m_i * m_right, s)
-
-                left_vecs_i = left_vecs_i.view(batch_size, m_left, m_i, m_right, s)
+                    left_vecs_i = left_vecs.view(batch_size, m_left_prev, m_left_i * m_left_next * s)
+                    left_vecs_i = _t_matmul(self.lazy_vars[i + 1 :], left_vecs_i)
+                    m_left_prev = m_right_prev
+                    left_vecs_i = left_vecs_i.contiguous().view(batch_size, m_left_prev, m_left_i * m_left_next, s)
+                else:
+                    left_vecs_i = left_vecs.view(batch_size, m_left_prev, m_left_i * m_left_next, s)
+
+                left_vecs_i = left_vecs_i.view(batch_size, m_left_prev, m_left_i, m_left_next, s)
                 left_vecs_i = left_vecs_i.transpose(1, 2)
                 if i != 0:
                     left_vecs_i = left_vecs_i.transpose(1, 3).contiguous()
-                    left_vecs_i = left_vecs_i.view(batch_size, m_right, m_left * m_i * s)
-                    left_vecs_i = _matmul(self.lazy_vars[:i], left_vecs_i)
-                    left_vecs_i = left_vecs_i.view(batch_size, m_right, m_left, m_i, s)
+                    left_vecs_i = left_vecs_i.view(batch_size, m_left_next, m_left_prev * m_left_i * s)
+                    left_vecs_i = _t_matmul(self.lazy_vars[:i], left_vecs_i)
+                    m_left_next = m_right_next
+                    left_vecs_i = left_vecs_i.view(batch_size, m_left_next, m_left_prev, m_left_i, s)
                     left_vecs_i = left_vecs_i.transpose(1, 3)
-                left_vecs_i = left_vecs_i.contiguous().view(batch_size, m_i, m_left * m_right * s)
+                left_vecs_i = left_vecs_i.contiguous().view(batch_size, m_left_i, m_left_prev * m_left_next * s)
                 left_vecs_i = left_vecs_i.contiguous()
                 res = res + list(lazy_var._quad_form_derivative(left_vecs_i, right_vecs_i))
 
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_variable.py b/gpytorch/lazy/lazy_evaluated_kernel_variable.py
index c03a77f2..e161fe48 100644
--- a/gpytorch/lazy/lazy_evaluated_kernel_variable.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_variable.py
@@ -37,6 +37,7 @@ class LazyEvaluatedKernelVariable(LazyVariable):
 
     def _get_indices(self, left_indices, right_indices):
         from ..kernels import Kernel
+
         x1 = self.x1[left_indices, :].unsqueeze(0)
         x2 = self.x2[right_indices, :].unsqueeze(0)
         res = super(Kernel, self.kernel).__call__(x1.transpose(0, 1), x2.transpose(0, 1))
diff --git a/gpytorch/likelihoods/softmax_likelihood.py b/gpytorch/likelihoods/softmax_likelihood.py
index 91606e03..10204cc7 100755
--- a/gpytorch/likelihoods/softmax_likelihood.py
+++ b/gpytorch/likelihoods/softmax_likelihood.py
@@ -38,9 +38,7 @@ class SoftmaxLikelihood(Likelihood):
             p(y|x) = \\Phi(\\frac{\\mu}{\\sqrt{1+\\sigma^2_f}})
         """
         if not isinstance(latent_func, GaussianRandomVariable):
-            raise RuntimeError(
-                "SoftmaxLikelihood expects a Gaussian distributed latent function to make predictions"
-            )
+            raise RuntimeError("SoftmaxLikelihood expects a Gaussian distributed latent function to make predictions")
 
         n_samples = settings.num_likelihood_samples.value()
         samples = latent_func.sample(n_samples)
diff --git a/gpytorch/utils/linear_cg.py b/gpytorch/utils/linear_cg.py
index 681a2fa4..b2fe04c7 100644
--- a/gpytorch/utils/linear_cg.py
+++ b/gpytorch/utils/linear_cg.py
@@ -179,7 +179,7 @@ def linear_cg(
         result = result.squeeze(-1)
 
     if n_tridiag:
-        t_mat = t_mat[:last_tridiag_iter + 1, :last_tridiag_iter + 1]
+        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:
diff --git a/test/lazy/test_kronecker_product_lazy_variable.py b/test/lazy/test_kronecker_product_lazy_variable.py
index f9cf4402..6db2c59f 100644
--- a/test/lazy/test_kronecker_product_lazy_variable.py
+++ b/test/lazy/test_kronecker_product_lazy_variable.py
@@ -57,27 +57,72 @@ class TestKroneckerProductLazyVariable(unittest.TestCase):
         self.assertTrue(approx_equal(cvar_copy.grad.data, cvar.grad.data))
         self.assertTrue(approx_equal(vec_copy.grad.data, vec.grad.data))
 
-    def test_matmul_mat(self):
-        avar = Variable(a, requires_grad=True)
-        bvar = Variable(b, requires_grad=True)
-        cvar = Variable(c, requires_grad=True)
-        mat = Variable(torch.randn(24, 5), requires_grad=True)
-        kp_lazy_var = KroneckerProductLazyVariable(NonLazyVariable(avar), NonLazyVariable(bvar), NonLazyVariable(cvar))
-        res = kp_lazy_var.matmul(mat)
+    def test_matmul_vec_random_rectangular(self):
+        ax = torch.randn(4, 2, 3)
+        bx = torch.randn(4, 5, 2)
+        cx = torch.randn(4, 6, 4)
+        rhsx = torch.randn(4, 3 * 2 * 4, 1)
+        rhsx = rhsx / torch.norm(rhsx)
+        ax_copy = Variable(ax, requires_grad=True)
+        bx_copy = bx.clone()
+        cx_copy = cx.clone()
+        rhsx_copy = rhsx.clone()
+
+        ax.requires_grad = True
+        bx.requires_grad = True
+        cx.requires_grad = True
+        ax_copy.requires_grad = True
+        bx_copy.requires_grad = True
+        cx_copy.requires_grad = True
+        rhsx.requires_grad = True
+        rhsx_copy.requires_grad = True
+
+        kp_lazy_var = KroneckerProductLazyVariable(NonLazyVariable(ax), NonLazyVariable(bx), NonLazyVariable(cx))
+        res = kp_lazy_var.matmul(rhsx)
+
+        actual_mat = kron(kron(ax_copy, bx_copy), cx_copy)
+        actual = actual_mat.matmul(rhsx_copy)
 
-        avar_copy = Variable(a, requires_grad=True)
-        bvar_copy = Variable(b, requires_grad=True)
-        cvar_copy = Variable(c, requires_grad=True)
-        mat_copy = Variable(mat.data.clone(), requires_grad=True)
-        actual = kron(kron(avar_copy, bvar_copy), cvar_copy).matmul(mat_copy)
         self.assertTrue(approx_equal(res.data, actual.data))
 
         actual.sum().backward()
         res.sum().backward()
-        self.assertTrue(approx_equal(avar_copy.grad.data, avar.grad.data))
-        self.assertTrue(approx_equal(bvar_copy.grad.data, bvar.grad.data))
-        self.assertTrue(approx_equal(cvar_copy.grad.data, cvar.grad.data))
-        self.assertTrue(approx_equal(mat_copy.grad.data, mat.grad.data))
+        self.assertTrue(approx_equal(ax_copy.grad.data, ax.grad.data))
+        self.assertTrue(approx_equal(bx_copy.grad.data, bx.grad.data))
+        self.assertTrue(approx_equal(cx_copy.grad.data, cx.grad.data))
+        self.assertTrue(approx_equal(rhsx_copy.grad.data, rhsx.grad.data))
+
+    def test_matmul_mat_random_rectangular(self):
+        a = torch.randn(4, 2, 3)
+        b = torch.randn(4, 5, 2)
+        c = torch.randn(4, 6, 4)
+        rhs = torch.randn(4, 3 * 2 * 4, 2)
+        a_copy = torch.tensor(a)
+        b_copy = b.clone()
+        c_copy = c.clone()
+        rhs_copy = rhs.clone()
+
+        a.requires_grad = True
+        b.requires_grad = True
+        c.requires_grad = True
+        a_copy.requires_grad = True
+        b_copy.requires_grad = True
+        c_copy.requires_grad = True
+        rhs.requires_grad = True
+        rhs_copy.requires_grad = True
+
+        actual = kron(kron(a_copy, b_copy), c_copy).matmul(rhs_copy)
+        kp_lazy_var = KroneckerProductLazyVariable(NonLazyVariable(a), NonLazyVariable(b), NonLazyVariable(c))
+        res = kp_lazy_var.matmul(rhs)
+
+        self.assertTrue(approx_equal(res.data, actual.data))
+
+        actual.sum().backward()
+        res.sum().backward()
+        self.assertTrue(approx_equal(a_copy.grad.data, a.grad.data))
+        self.assertTrue(approx_equal(b_copy.grad.data, b.grad.data))
+        self.assertTrue(approx_equal(c_copy.grad.data, c.grad.data))
+        self.assertTrue(approx_equal(rhs_copy.grad.data, rhs.grad.data))
 
     def test_matmul_batch_mat(self):
         avar = Variable(a.repeat(3, 1, 1), requires_grad=True)
diff --git a/test/util/test_linear_cg.py b/test/util/test_linear_cg.py
index e9806b1d..cd98f9db 100644
--- a/test/util/test_linear_cg.py
+++ b/test/util/test_linear_cg.py
@@ -14,11 +14,11 @@ from gpytorch.utils.linear_cg import linear_cg
 class TestLinearCG(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(0)
-                if torch.cuda.is_available():
-                    torch.cuda.manual_seed_all(0)
-                random.seed(0)
+            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)
 
     def tearDown(self):
         if hasattr(self, "rng_state"):