8c40ed8fb56baef7ab5ce8d5ed85bb10e4f09174,gpytorch/lazy/kronecker_product_lazy_tensor.py,KroneckerProductLazyTensor,root_decomposition,#KroneckerProductLazyTensor#Any#,199

Before Change

    def root_decomposition(self, method: Optional[str] = None):
        from gpytorch.lazy import RootLazyTensor

        if method == "symeig" or method is None:
            evals, evecs = self._symeig(eigenvectors=True, return_evals_as_lazy=True)
            // TODO: only use non-zero evals (req. dealing w/ batches...)
            f_list = [
                evec * eval.diag().clamp(0.0).sqrt().unsqueeze(-2)
                for eval, evec in zip(evals.lazy_tensors, evecs.lazy_tensors)
            ]
            F = KroneckerProductLazyTensor(*f_list)
            return RootLazyTensor(F)
        else:
            return super().root_decomposition(method=method)

    @cached(name="size")
    def _size(self):
        left_size = _prod(lazy_tensor.size(-2) for lazy_tensor in self.lazy_tensors)
        right_size = _prod(lazy_tensor.size(-1) for lazy_tensor in self.lazy_tensors)

After Change

        from gpytorch.lazy import RootLazyTensor

        // return a dense root decomposition if the matrix is small
        if self.shape[-1] <= settings.max_cholesky_size.value():
            return super().root_decomposition(method=method)

        root_list = [lt.root_decomposition(method=method).root for lt in self.lazy_tensors]
        kronecker_root = KroneckerProductLazyTensor(*root_list)
        return RootLazyTensor(kronecker_root)

Instances