ad18add8ae7da04213813b5675b903bbf4be1ebd,gpytorch/lazy/kronecker_product_added_diag_lazy_tensor.py,KroneckerProductAddedDiagLazyTensor,_root_inv_decomposition,#KroneckerProductAddedDiagLazyTensor#Any#,219

Before Change

        lt = self.lazy_tensor
        if isinstance(self.diag_tensor, KroneckerProductDiagLazyTensor):
            if all(isinstance(tdiag, ConstantDiagLazyTensor) for tdiag in dlt.lazy_tensors):
                sub_evals, sub_evecs = [], []
                for lt_, dlt_ in zip(lt.lazy_tensors, dlt.lazy_tensors):
                    evals_, evecs_ = lt_.symeig(eigenvectors=True)
                    sub_evals.append(DiagLazyTensor(evals_ / dlt_.diag_values))
                    sub_evecs.append(evecs_)

After Change

                // so that we compute
                // L^{-1/2} = D^{1/2} Q (\kron a_i^{-1} \Lambda_i + I)^{1/2}
                //       = (\kron a_i^{1/2} Q_i)(\kron a_i^{-1} \Lambda_i + I)^{-1/2}
                scaled_evecs_list = []
                for evec_, dlt_ in zip(evecs.lazy_tensors, dlt.lazy_tensors):
                    scaled_evecs_list.append(evec_ * dlt_.diag_values.sqrt())
                scaled_evecs = KroneckerProductLazyTensor(*scaled_evecs_list)
                return MatmulLazyTensor(scaled_evecs, evals_p_i_inv_root)

            // again, we compute the root decomposition by pulling across the diagonals
            dlt_sqrt, evals_p_i, evecs = _symmetrize_kpadlt_constructor(lt, dlt)

Instances