Merge pull request #2388 from douglas-boubert/grad_kernels_only_create_covariance_if_needed

Stop rbf_kernel_grad and rbf_kernel_gradgrad creating the full covariance matrix unnecessarily

Author: Balandat

gpytorch/kernels/rbf_kernel_grad.py

@@ -57,9 +57,9 @@ def forward(self, x1, x2, diag=False, **params):
         n1, d = x1.shape[-2:]
         n2 = x2.shape[-2]
 
-        K = torch.zeros(*batch_shape, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype)
-
         if not diag:
+            K = torch.zeros(*batch_shape, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype)
+
             # Scale the inputs by the lengthscale (for stability)
             x1_ = x1.div(self.lengthscale)
             x2_ = x2.div(self.lengthscale)

gpytorch/kernels/rbf_kernel_gradgrad.py

@@ -57,9 +57,9 @@ def forward(self, x1, x2, diag=False, **params):
         n1, d = x1.shape[-2:]
         n2 = x2.shape[-2]
 
-        K = torch.zeros(*batch_shape, n1 * (2 * d + 1), n2 * (2 * d + 1), device=x1.device, dtype=x1.dtype)
-
         if not diag:
+            K = torch.zeros(*batch_shape, n1 * (2 * d + 1), n2 * (2 * d + 1), device=x1.device, dtype=x1.dtype)
+
             # Scale the inputs by the lengthscale (for stability)
             x1_ = x1.div(self.lengthscale)
             x2_ = x2.div(self.lengthscale)