Adds settings.trace_mode and skip cholesky jitter

linear_operator/settings.py

@@ -551,6 +551,24 @@ class terminate_cg_by_size(_feature_flag):
     _default = False
 
 
+class trace_mode(_feature_flag):
+    """
+    If set to True, we will generally try to avoid calling our built in PyTorch functions, because these cannot
+    be run through torch.jit.trace.
+
+    Note that this will sometimes involve explicitly evaluating lazy tensors and various other slowdowns and
+    inefficiencies. As a result, you really shouldn't use this feature context unless you are calling torch.jit.trace
+    on a GPyTorch model.
+
+    Our hope is that this flag will not be necessary long term, once https://github.com/pytorch/pytorch/issues/22329
+    is fixed.
+
+    (Default: False)
+    """
+
+    _default = False
+
+
 class tridiagonal_jitter(_value_context):
     """
     The (relative) amount of noise to add to the diagonal of tridiagonal matrices before

linear_operator/utils/cholesky.py

@@ -18,7 +18,7 @@ def _psd_safe_cholesky(A, out=None, jitter=None, max_tries=None):
         out = (out, torch.empty(A.shape[:-2], dtype=torch.int32, device=out.device))
 
     L, info = torch.linalg.cholesky_ex(A, out=out)
-    if not torch.any(info):
+    if settings.trace_mode.on() or not torch.any(info):
         return L
 
     isnan = torch.isnan(A)