evals_, evecs_ = lt_.symeig(eigenvectors=True)
sub_evals.append(DiagLazyTensor(evals_ / dlt_.diag_values))
sub_evecs.append(evecs_)
Lambda_I = KroneckerProductDiagLazyTensor(*sub_evals).add_jitter(1.0)
S = KroneckerProductLazyTensor(*sub_evecs)
return MatmulLazyTensor(S, Lambda_I.sqrt())
else:
// again, we compute the root decomposition by pulling across the diagonals
After Change
// so that we compute
// L = 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_root)
// again, we compute the root decomposition by pulling across the diagonals
dlt_root = dlt.sqrt()