e8a945606e0b3dafe71ce2f1bf1b52ab3b05428c,gpytorch/lazy/kronecker_product_added_diag_lazy_tensor.py,KroneckerProductAddedDiagLazyTensor,_root_decomposition,#KroneckerProductAddedDiagLazyTensor#,186

Before Change

                return MatmulLazyTensor(S, Lambda_I.sqrt())
            else:    
                // again, we compute the root decomposition by pulling across the diagonals
                dlt_inv_root = dlt.sqrt().inverse()
                dlt_root = dlt.sqrt()
                symm_prod = KroneckerProductLazyTensor(
                    *[d.matmul(k).matmul(d) for k, d in zip(lt.lazy_tensors, dlt_inv_root.lazy_tensors)]
                )
                evals, evecs = symm_prod.symeig(eigenvectors = True)
                evals_p_i_root = DiagLazyTensor((evals + 1.).sqrt())
                return MatmulLazyTensor(dlt_root, MatmulLazyTensor(evecs, evals_p_i_root))
            
        return super()._root_decomposition()

After Change

            else:
                // again, we compute the root decomposition by pulling across the diagonals
                dlt_root = dlt.sqrt()
                _, evals_p_i, evecs = _symmetrize_kpadlt_constructor(lt, dlt)
                evals_p_i_root = DiagLazyTensor(evals_p_i.diag().sqrt())
                return MatmulLazyTensor(dlt_root, MatmulLazyTensor(evecs, evals_p_i_root))

        return super()._root_decomposition()

Instances