add a check for root existing in iqlt
@@ -23,7 +23,7 @@ from ..utils.cholesky import psd_safe_cholesky
from ..utils.deprecation import _deprecate_renamed_methods
from ..utils.errors import CachingError
from ..utils.getitem import _compute_getitem_size, _convert_indices_to_tensors, _is_noop_index, _noop_index
-from ..utils.memoize import add_to_cache, cached, get_from_cache, pop_from_cache
+from ..utils.memoize import _is_in_cache_ignore_args, add_to_cache, cached, get_from_cache, pop_from_cache
from ..utils.pivoted_cholesky import pivoted_cholesky
from ..utils.warnings import NumericalWarning
from .lazy_tensor_representation_tree import LazyTensorRepresentationTree
@@ -1028,6 +1028,13 @@ class LazyTensor(ABC):
from .chol_lazy_tensor import CholLazyTensor
from .triangular_lazy_tensor import TriangularLazyTensor
+ # if the root decomposition has already been computed and is triangular we can use it instead
+ # of computing the cholesky.
+ if _is_in_cache_ignore_args(self, "root_decomposition"):
+ root = self.root_decomposition().root
+ if isinstance(root, TriangularLazyTensor):
+ cholesky = root
+ else:
cholesky = CholLazyTensor(TriangularLazyTensor(self.cholesky()))
return cholesky.inv_quad_logdet(inv_quad_rhs=inv_quad_rhs, logdet=logdet, reduce_inv_quad=reduce_inv_quad)
