Repository: cornellius-gp/gpytorch

Issue:
MultitaskGPModel won't scale with size of input
In the basic multi-task model example shown in https://nbviewer.jupyter.org/github/cornellius-gp/gpytorch/blob/master/examples/multitask_gp_regression.ipynb, the code doesn't scale to larger sizes of training data. Specifically, if you changed training data size in cell 2 of the notebook from 11 to 100, the code no longer works. It throws the following error:
RuntimeError: Function InvQuadLogDetLegacyBackward returned an invalid gradient at index 1 - expected shape [200, 9] but got [200, 5]

Can someone figure out what's going on? Or am I doing this wrong?


Discussion:
I'm looking in to this -- it has to do with the recent changes to lazily evaluating kernels. Easy enough to fix in the exact GP case (I'll push something soon), a little trickier in the sparse GP setting. Working on a fix now.

Resolving patch:
diff --git a/gpytorch/lazy/mul_lazy_variable.py b/gpytorch/lazy/mul_lazy_variable.py
--- a/gpytorch/lazy/mul_lazy_variable.py
+++ b/gpytorch/lazy/mul_lazy_variable.py
@@ -3,6 +3,7 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import torch
 from torch.autograd import Variable
 from .lazy_variable import LazyVariable
 from .non_lazy_variable import NonLazyVariable
@@ -18,18 +19,30 @@ def __init__(self, *lazy_vars):
         """
         lazy_vars = list(lazy_vars)
         if len(lazy_vars) == 1:
-            raise RuntimeError("MulLazyVariable should have more than one lazy variables")
-
-        for i, lazy_var in enumerate(lazy_vars):
-            if not isinstance(lazy_var, LazyVariable):
-                if isinstance(lazy_var, Variable):
-                    lazy_vars[i] = NonLazyVariable(lazy_var)
-                else:
-                    raise RuntimeError("All arguments of a MulLazyVariable should be lazy " "variables or variables")
+            if isinstance(lazy_vars[0], NonLazyVariable):
+                self._non_lazy_self = lazy_vars
+            else:
+                raise RuntimeError("MulLazyVariable can only have one LazyVariable if it is a NonLazyVariable")
+        else:
+            for i, lazy_var in enumerate(lazy_vars):
+                if not isinstance(lazy_var, LazyVariable):
+                    if torch.is_tensor(lazy_var):
+                        lazy_vars[i] = NonLazyVariable(lazy_var)
+                    else:
+                        raise RuntimeError("All arguments of a MulLazyVariable should be lazy variables or variables")
 
         super(MulLazyVariable, self).__init__(*lazy_vars)
         self.lazy_vars = lazy_vars
 
+    @property
+    def non_lazy_self(self):
+        if hasattr(self, "_non_lazy_self"):
+            return self._non_lazy_self[0]
+        elif len(self._args) == 1:
+            return self._args[0]
+        else:
+            return None
+
     @property
     def left_lazy_var(self):
         return self._args[0]
@@ -40,44 +53,53 @@ def right_lazy_var(self):
 
     @property
     def _args(self):
-        if not hasattr(self, "_mul_args_memo"):
-            # Sort lazy variables by root decomposition size (rank)
-            lazy_vars = sorted(self.lazy_vars, key=lambda lv: lv.root_decomposition_size())
-
-            # Recursively construct lazy variables
-            # Make sure the recursive components get a mix of low_rank and high_rank variables
-            if len(lazy_vars) > 2:
-                interleaved_lazy_vars = lazy_vars[0::2] + lazy_vars[1::2]
-                if len(interleaved_lazy_vars) > 3:
-                    left_lazy_var = MulLazyVariable(*interleaved_lazy_vars[: len(interleaved_lazy_vars) // 2])
-                    if left_lazy_var.root_decomposition_size() < left_lazy_var.size(-1):
-                        left_lazy_var = RootLazyVariable(left_lazy_var.root_decomposition())
+        if not hasattr(self, "_mul_args_memo") and not hasattr(self, "_non_lazy_self"):
+            lazy_vars = sorted(
+                (lv.evaluate_kernel() for lv in self.lazy_vars), key=lambda lv: lv.root_decomposition_size()
+            )
+
+            if any(isinstance(lv, NonLazyVariable) for lv in lazy_vars):
+                self._non_lazy_self = [NonLazyVariable(prod([lv.evaluate() for lv in lazy_vars]))]
+            else:
+                # Sort lazy variables by root decomposition size (rank)
+
+                # Recursively construct lazy variables
+                # Make sure the recursive components get a mix of low_rank and high_rank variables
+                if len(lazy_vars) > 2:
+                    interleaved_lazy_vars = lazy_vars[0::2] + lazy_vars[1::2]
+                    if len(interleaved_lazy_vars) > 3:
+                        left_lazy_var = MulLazyVariable(*interleaved_lazy_vars[: len(interleaved_lazy_vars) // 2])
+                        if left_lazy_var.root_decomposition_size() < left_lazy_var.size(-1):
+                            left_lazy_var = RootLazyVariable(left_lazy_var.root_decomposition())
+                        else:
+                            left_lazy_var = NonLazyVariable(left_lazy_var.evaluate())
                     else:
-                        left_lazy_var = NonLazyVariable(left_lazy_var.evaluate())
-                else:
-                    # Make sure we're not constructing a MulLazyVariable of length 1
-                    left_lazy_var = interleaved_lazy_vars[0]
+                        # Make sure we're not constructing a MulLazyVariable of length 1
+                        left_lazy_var = interleaved_lazy_vars[0]
 
-                right_lazy_var = MulLazyVariable(*interleaved_lazy_vars[len(interleaved_lazy_vars) // 2 :])
-                if right_lazy_var.root_decomposition_size() < right_lazy_var.size(-1):
-                    right_lazy_var = RootLazyVariable(right_lazy_var.root_decomposition())
+                    right_lazy_var = MulLazyVariable(*interleaved_lazy_vars[len(interleaved_lazy_vars) // 2 :])
+                    if right_lazy_var.root_decomposition_size() < right_lazy_var.size(-1):
+                        right_lazy_var = RootLazyVariable(right_lazy_var.root_decomposition())
+                    else:
+                        right_lazy_var = NonLazyVariable(right_lazy_var.evaluate())
                 else:
-                    right_lazy_var = NonLazyVariable(right_lazy_var.evaluate())
-            else:
-                left_lazy_var = lazy_vars[0]
-                right_lazy_var = lazy_vars[1]
+                    left_lazy_var = lazy_vars[0]
+                    right_lazy_var = lazy_vars[1]
 
-            # Choose which we're doing: root decomposition or exact
-            if left_lazy_var.root_decomposition_size() < left_lazy_var.size(-1):
-                left_lazy_var = RootLazyVariable(left_lazy_var.root_decomposition())
-                right_lazy_var = RootLazyVariable(right_lazy_var.root_decomposition())
-            else:
-                left_lazy_var = NonLazyVariable(left_lazy_var.evaluate())
-                right_lazy_var = NonLazyVariable(right_lazy_var.evaluate())
+                # Choose which we're doing: root decomposition or exact
+                if left_lazy_var.root_decomposition_size() < left_lazy_var.size(-1):
+                    left_lazy_var = RootLazyVariable(left_lazy_var.root_decomposition())
+                    right_lazy_var = RootLazyVariable(right_lazy_var.root_decomposition())
 
-            self._mul_args_memo = [left_lazy_var, right_lazy_var]
+                if isinstance(left_lazy_var, NonLazyVariable) and isinstance(right_lazy_var, NonLazyVariable):
+                    self._non_lazy_self = [NonLazyVariable(left_lazy_var.evaluate() * right_lazy_var.evaluate())]
+                else:
+                    self._mul_args_memo = [left_lazy_var, right_lazy_var]
 
-        return self._mul_args_memo
+        if hasattr(self, "_mul_args_memo"):
+            return self._mul_args_memo
+        else:
+            return self._non_lazy_self
 
     @_args.setter
     def _args(self, args):
@@ -85,6 +107,9 @@ def _args(self, args):
         pass
 
     def _matmul(self, rhs):
+        if self.non_lazy_self is not None:
+            return self.non_lazy_self._matmul(rhs)
+
         is_vector = False
         if rhs.ndimension() == 1:
             rhs = rhs.unsqueeze(1)
@@ -110,6 +135,9 @@ def _matmul(self, rhs):
         return res
 
     def _quad_form_derivative(self, left_vecs, right_vecs):
+        if self.non_lazy_self is not None:
+            return self.non_lazy_self._quad_form_derivative(left_vecs, right_vecs)
+
         if left_vecs.ndimension() == 1:
             left_vecs = left_vecs.unsqueeze(1)
             right_vecs = right_vecs.unsqueeze(1)
@@ -178,6 +206,10 @@ def mul(self, other):
         elif isinstance(other, MulLazyVariable):
             res = list(self.lazy_vars) + list(other.lazy_vars)
             return MulLazyVariable(*res)
+        elif isinstance(other, NonLazyVariable):
+            return NonLazyVariable(self.evaluate() * other.evaluate())
+        elif self.non_lazy_self is not None:
+            return NonLazyVariable(self.non_lazy_self.evaluate() * other.evaluate())
         elif isinstance(other, LazyVariable):
             return MulLazyVariable(*(list(self.lazy_vars) + [other]))
         else:
@@ -199,3 +231,18 @@ def _get_indices(self, left_indices, right_indices):
     def _transpose_nonbatch(self):
         # mul_lazy_variable only works with symmetric matrices
         return self
+
+    def representation(self):
+        """
+        Returns the variables that are used to define the LazyVariable
+        """
+        if self.non_lazy_self is not None:
+            return self.non_lazy_self.representation()
+        else:
+            return super(MulLazyVariable, self).representation()
+
+    def representation_tree(self):
+        if self.non_lazy_self is not None:
+            return self.non_lazy_self.representation_tree()
+        else:
+            return super(MulLazyVariable, self).representation_tree()


Test patch:
diff --git a/test/examples/test_multitask_gp_regression.py b/test/examples/test_multitask_gp_regression.py
--- a/test/examples/test_multitask_gp_regression.py
+++ b/test/examples/test_multitask_gp_regression.py
@@ -19,9 +19,9 @@
 from gpytorch.random_variables import GaussianRandomVariable
 
 # Simple training data: let's try to learn a sine function
-train_x = Variable(torch.linspace(0, 1, 11))
-y1_inds = Variable(torch.zeros(11).long())
-y2_inds = Variable(torch.ones(11).long())
+train_x = Variable(torch.linspace(0, 1, 100))
+y1_inds = Variable(torch.zeros(100).long())
+y2_inds = Variable(torch.ones(100).long())
 train_y1 = Variable(torch.sin(train_x.data * (2 * pi)))
 train_y2 = Variable(torch.cos(train_x.data * (2 * pi)))