fix keops bypass in rbf kernel
@@ -26,12 +26,16 @@ try:
)
def covar_func(self, x1, x2, diag=False):
- # TODO: x1 / x2 size checks are a work around for a very minor bug in KeOps.
- # This bug is fixed on KeOps master, and we'll remove that part of the check
- # when they cut a new release.
- elif x1.size(-2) == 1 or x2.size(-2) == 1:
+ # We only should use KeOps on big kernel matrices
+ # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)
+ # then don't apply KeOps
+ if (
+ diag
+ or x1.size(-2) < settings.max_cholesky_size.value()
+ or x2.size(-2) < settings.max_cholesky_size.value()
+ ):
return self._nonkeops_covar_func(x1, x2, diag=diag)
- else:
+
with torch.autograd.enable_grad():
x1_ = KEOLazyTensor(x1[..., :, None, :])
x2_ = KEOLazyTensor(x2[..., None, :, :])
@@ -44,17 +48,11 @@ try:
x1_ = x1.div(self.lengthscale)
x2_ = x2.div(self.lengthscale)
- # We only should use KeOps on big kernel matrices
- # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)
- # then don't apply KeOps
- if (
- diag
- or x1_.size(-2) < settings.max_cholesky_size.value()
- or x2_.size(-2) < settings.max_cholesky_size.value()
- ):
- return self._nonkeops_covar_func(x1_, x2_, diag=diag)
+ covar_func = lambda x1, x2, diag=diag: self.covar_func(x1, x2, diag)
+
+ if diag:
+ return covar_func(x1_, x2_, diag=True)
- covar_func = lambda x1, x2, diag=False: self.covar_func(x1, x2, diag)
return KeOpsLazyTensor(x1_, x2_, covar_func)
