jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def compute_sparse_code_FISTA(input_batch, dictionary, lambda_star, maxiter, tol=1e-3):
    with torch.no_grad():
        L = torch.symeig(torch.mm(dictionary[..., 0, 0].t(), dictionary[..., 0, 0]))[0][-1]
        input_size, dict_size = dictionary.size(0), dictionary.size(1)

        support_size_FISTA_curve = np.zeros(maxiter)

        batch_size, M, N = input_batch.size(0), input_batch.size(2), input_batch.size(3)
        x = input_batch.new_zeros(batch_size, dict_size, M, N)
        t = 1
        y = x.clone()
        for i_iter in range(maxiter):
            x_old = x.clone()
            D_y = nn.functional.conv2d(y, dictionary)
            y = y + nn.functional.conv2d(input_batch - D_y, dictionary.transpose(0, 1).contiguous()) / L
            x = relu(y, lambda_star / L)
            t0 = t
            t = (1. + math.sqrt(1. + 4. * t ** 2)) / 2.
            y = x + ((t0 - 1.) / t) * (x - x_old)

            support_size_FISTA_curve[i_iter] = (x != 0).sum().item()

            diff = ((x - x_old).norm(p=2, dim=1) / (x.norm(p=2, dim=1) + 1e-10))

            if tol is not None and diff.max() < tol:
                support_size_FISTA_curve[i_iter:] = support_size_FISTA_curve[i_iter - 1]
                break

    return x, support_size_FISTA_curve
//  [python]
def compute_sparse_code_FISTA(input_batch, dictionary, lambda_star, maxiter, tol=1e-3):
    with torch.no_grad():
        L = torch.linalg.eigh(torch.mm(dictionary[..., 0, 0].t(), dictionary[..., 0, 0]))[0][-1]
        input_size, dict_size = dictionary.size(0), dictionary.size(1)

        support_size_FISTA_curve = np.zeros(maxiter)

        batch_size, M, N = input_batch.size(0), input_batch.size(2), input_batch.size(3)
        x = input_batch.new_zeros(batch_size, dict_size, M, N)
        t = 1
        y = x.clone()
        for i_iter in range(maxiter):
            x_old = x.clone()
            D_y = nn.functional.conv2d(y, dictionary)
            y = y + nn.functional.conv2d(input_batch - D_y, dictionary.transpose(0, 1).contiguous()) / L
            x = relu(y, lambda_star / L)
            t0 = t
            t = (1. + math.sqrt(1. + 4. * t ** 2)) / 2.
            y = x + ((t0 - 1.) / t) * (x - x_old)

            support_size_FISTA_curve[i_iter] = (x != 0).sum().item()

            diff = ((x - x_old).norm(p=2, dim=1) / (x.norm(p=2, dim=1) + 1e-10))

            if tol is not None and diff.max() < tol:
                support_size_FISTA_curve[i_iter:] = support_size_FISTA_curve[i_iter - 1]
                break

    return x, support_size_FISTA_curve
//  [python]
def compute_sparse_code_ISTA(input_batch, dictionary, lambda_star, maxiter, tol=1e-3):
    with torch.no_grad():
        L = torch.symeig(torch.mm(dictionary[..., 0, 0].t(), dictionary[..., 0, 0]))[0][-1]
        input_size, dict_size = dictionary.size(0), dictionary.size(1)

        support_size_ISTA_curve = np.zeros(maxiter)

        batch_size, M, N = input_batch.size(0), input_batch.size(2), input_batch.size(3)
        x = input_batch.new_zeros(batch_size, dict_size, M, N)
        for i_iter in range(maxiter):
            x_old = x.clone()
            D_x = nn.functional.conv2d(x, dictionary)
            x = x + nn.functional.conv2d(input_batch - D_x, dictionary.transpose(0, 1).contiguous()) / L
            x = relu(x, lambda_star / L)

            support_size_ISTA_curve[i_iter] = (x != 0).sum().item()

            diff = ((x - x_old).norm(p=2, dim=1) / (x.norm(p=2, dim=1) + 1e-10))

            if tol is not None and diff.max() < tol:
                support_size_ISTA_curve[i_iter:] = support_size_ISTA_curve[i_iter - 1]
                break

    return x, support_size_ISTA_curve
//  [python]
def compute_sparse_code_ISTA(input_batch, dictionary, lambda_star, maxiter, tol=1e-3):
    with torch.no_grad():
        L = torch.linalg.eigh(torch.mm(dictionary[..., 0, 0].t(), dictionary[..., 0, 0]))[0][-1]
        input_size, dict_size = dictionary.size(0), dictionary.size(1)

        support_size_ISTA_curve = np.zeros(maxiter)

        batch_size, M, N = input_batch.size(0), input_batch.size(2), input_batch.size(3)
        x = input_batch.new_zeros(batch_size, dict_size, M, N)
        for i_iter in range(maxiter):
            x_old = x.clone()
            D_x = nn.functional.conv2d(x, dictionary)
            x = x + nn.functional.conv2d(input_batch - D_x, dictionary.transpose(0, 1).contiguous()) / L
            x = relu(x, lambda_star / L)

            support_size_ISTA_curve[i_iter] = (x != 0).sum().item()

            diff = ((x - x_old).norm(p=2, dim=1) / (x.norm(p=2, dim=1) + 1e-10))

            if tol is not None and diff.max() < tol:
                support_size_ISTA_curve[i_iter:] = support_size_ISTA_curve[i_iter - 1]
                break

    return x, support_size_ISTA_curve
//  [python]
def test_fractional_matrix_power(N0=5):
    np0 = rand_hermite_matrix(N0, tag_complex=False)
    np1 = np.random.rand()
    ret_ = scipy.linalg.fractional_matrix_power(np0, np1)
    EVL,EVC = torch.symeig(torch.tensor(np0), eigenvectors=True)
    ret0 = torch.matmul(EVC * (EVL**np1), EVC.t_()).numpy()
    assert hfe(ret_, ret0) < 1e-7
//  [python]
def test_fractional_matrix_power(N0=5):
    np0 = rand_hermite_matrix(N0, tag_complex=False)
    np1 = np.random.rand()
    ret_ = scipy.linalg.fractional_matrix_power(np0, np1)
    EVL,EVC = torch.linalg.eigh(torch.tensor(np0))
    ret0 = torch.matmul(EVC * (EVL**np1), EVC.t_()).numpy()
    assert hfe(ret_, ret0) < 1e-7
//  [python]
def test_torch_symeig(N0=3, N1=23):
    tmp0 = np.random.rand(N0, N1, N1)
    np0 = tmp0 + tmp0.transpose(0,2,1)
    ret_ = np.stack([np.linalg.eigvalsh(x) for x in np0])
    EVL,EVC = torch.symeig(torch.tensor(np0), eigenvectors=True)
    EVL, EVC = EVL.numpy(), EVC.numpy()
    assert hfe(ret_, EVL) < 1e-7
    assert max(hfe(np.eye(N1), x@x.T) for x in EVC) < 1e-7
    assert hfe(np0, np.stack([(y*x)@y.T for x,y in zip(EVL,EVC)])) < 1e-7
//  [python]
def test_torch_symeig(N0=3, N1=23):
    tmp0 = np.random.rand(N0, N1, N1)
    np0 = tmp0 + tmp0.transpose(0,2,1)
    ret_ = np.stack([np.linalg.eigvalsh(x) for x in np0])
    EVL,EVC = torch.linalg.eigh(torch.tensor(np0))
    EVL, EVC = EVL.numpy(), EVC.numpy()
    assert hfe(ret_, EVL) < 1e-7
    assert max(hfe(np.eye(N1), x@x.T) for x in EVC) < 1e-7
    assert hfe(np0, np.stack([(y*x)@y.T for x,y in zip(EVL,EVC)])) < 1e-7
//  [python]
def compute_conv_ISTA(model, val_loader, dictionary, nb_iter, override_lambda_star=None):
    with torch.no_grad():
        L = torch.symeig(torch.mm(dictionary[..., 0, 0].t(), dictionary[..., 0, 0]))[0][-1]
        input_size, dict_size = dictionary.size(0), dictionary.size(1)

        if override_lambda_star is not None:
            lambda_star = override_lambda_star
        else:
            lambda_star = model.module.istc.lambda_star

        loss_curve = np.zeros(nb_iter)
        conv = np.zeros(nb_iter)
        x_star_norm = np.zeros(1)

        for _, (input_batch, target) in enumerate(val_loader):
            input_batch = input_batch.cuda()
            input_batch = model(input_batch, return_proj=True)
            batch_size, M, N = input_batch.size(0), input_batch.size(2), input_batch.size(3)
            # Use FISTA to compute the l1 problem solution
            x_star = compute_sparse_code_FISTA(input_batch, dictionary, lambda_star, maxiter=200, tol=1e-3)[0]
            x_star_norm += (x_star.norm(p=2) ** 2).sum().item()

            x = input_batch.new_zeros(batch_size, dict_size, M, N)
            for i_iter in range(nb_iter):
                D_x = nn.functional.conv2d(x, dictionary)
                x = x + nn.functional.conv2d(input_batch - D_x, dictionary.transpose(0, 1).contiguous()) / L
                x = relu(x, lambda_star / L)

                rec_error = 0.5 * ((nn.functional.conv2d(x, dictionary) - input_batch).norm(p=2, dim=1) ** 2).sum().item()
                sparsity_loss = (lambda_star * x).norm(p=1, dim=1).sum().item()
                loss_curve[i_iter] += rec_error + sparsity_loss
                conv[i_iter] += ((x - x_star).norm(p=2) ** 2).sum().item()

        loss_curve /= (M * N * len(val_loader.dataset))
        conv_rel = conv / x_star_norm

    return loss_curve, conv_rel
//  [python]
def compute_conv_ISTA(model, val_loader, dictionary, nb_iter, override_lambda_star=None):
    with torch.no_grad():
        L = torch.linalg.eigh(torch.mm(dictionary[..., 0, 0].t(), dictionary[..., 0, 0]))[0][-1]
        input_size, dict_size = dictionary.size(0), dictionary.size(1)

        if override_lambda_star is not None:
            lambda_star = override_lambda_star
        else:
            lambda_star = model.module.istc.lambda_star

        loss_curve = np.zeros(nb_iter)
        conv = np.zeros(nb_iter)
        x_star_norm = np.zeros(1)

        for _, (input_batch, target) in enumerate(val_loader):
            input_batch = input_batch.cuda()
            input_batch = model(input_batch, return_proj=True)
            batch_size, M, N = input_batch.size(0), input_batch.size(2), input_batch.size(3)
            # Use FISTA to compute the l1 problem solution
            x_star = compute_sparse_code_FISTA(input_batch, dictionary, lambda_star, maxiter=200, tol=1e-3)[0]
            x_star_norm += (x_star.norm(p=2) ** 2).sum().item()

            x = input_batch.new_zeros(batch_size, dict_size, M, N)
            for i_iter in range(nb_iter):
                D_x = nn.functional.conv2d(x, dictionary)
                x = x + nn.functional.conv2d(input_batch - D_x, dictionary.transpose(0, 1).contiguous()) / L
                x = relu(x, lambda_star / L)

                rec_error = 0.5 * ((nn.functional.conv2d(x, dictionary) - input_batch).norm(p=2, dim=1) ** 2).sum().item()
                sparsity_loss = (lambda_star * x).norm(p=1, dim=1).sum().item()
                loss_curve[i_iter] += rec_error + sparsity_loss
                conv[i_iter] += ((x - x_star).norm(p=2) ** 2).sum().item()

        loss_curve /= (M * N * len(val_loader.dataset))
        conv_rel = conv / x_star_norm

    return loss_curve, conv_rel
//  [python]
def compute_conv_FISTA(model, val_loader, dictionary, nb_iter, override_lambda_star=None):
    with torch.no_grad():
        L = torch.symeig(torch.mm(dictionary[..., 0, 0].t(), dictionary[..., 0, 0]))[0][-1]
        input_size, dict_size = dictionary.size(0), dictionary.size(1)

        if override_lambda_star is not None:
            lambda_star = override_lambda_star
        else:
            lambda_star = model.module.istc.lambda_star

        loss_curve = np.zeros(nb_iter)
        conv = np.zeros(nb_iter)
        x_star_norm = np.zeros(1)

        for _, (input_batch, target) in enumerate(val_loader):
            input_batch = input_batch.cuda()
            input_batch = model(input_batch, return_proj=True)
            batch_size, M, N = input_batch.size(0), input_batch.size(2), input_batch.size(3)
            # Use FISTA to compute the l1 problem solution
            x_star = compute_sparse_code_FISTA(input_batch, dictionary, lambda_star, maxiter=200, tol=1e-3)[0]
            x_star_norm += (x_star.norm(p=2) ** 2).sum().item()

            x = input_batch.new_zeros(batch_size, dict_size, M, N)
            t = 1
            y = x.clone()
            for i_iter in range(nb_iter):
                x_old = x.clone()
                D_y = nn.functional.conv2d(y, dictionary)
                y = y + nn.functional.conv2d(input_batch - D_y, dictionary.transpose(0, 1).contiguous()) / L
                x = relu(y, lambda_star / L)
                t0 = t
                t = (1. + math.sqrt(1. + 4. * t ** 2)) / 2.
                y = x + ((t0 - 1.) / t) * (x - x_old)

                rec_error = 0.5 * ((nn.functional.conv2d(x, dictionary) - input_batch).norm(p=2, dim=1) ** 2).sum().item()
                sparsity_loss = (lambda_star * x).norm(p=1, dim=1).sum().item()
                loss_curve[i_iter] += rec_error + sparsity_loss
                conv[i_iter] += ((x - x_star).norm(p=2) ** 2).sum().item()

        loss_curve /= (M * N * len(val_loader.dataset))
        conv_rel = conv / x_star_norm

    return loss_curve, conv_rel
//  [python]
def compute_conv_FISTA(model, val_loader, dictionary, nb_iter, override_lambda_star=None):
    with torch.no_grad():
        L = torch.linalg.eigh(torch.mm(dictionary[..., 0, 0].t(), dictionary[..., 0, 0]))[0][-1]
        input_size, dict_size = dictionary.size(0), dictionary.size(1)

        if override_lambda_star is not None:
            lambda_star = override_lambda_star
        else:
            lambda_star = model.module.istc.lambda_star

        loss_curve = np.zeros(nb_iter)
        conv = np.zeros(nb_iter)
        x_star_norm = np.zeros(1)

        for _, (input_batch, target) in enumerate(val_loader):
            input_batch = input_batch.cuda()
            input_batch = model(input_batch, return_proj=True)
            batch_size, M, N = input_batch.size(0), input_batch.size(2), input_batch.size(3)
            # Use FISTA to compute the l1 problem solution
            x_star = compute_sparse_code_FISTA(input_batch, dictionary, lambda_star, maxiter=200, tol=1e-3)[0]
            x_star_norm += (x_star.norm(p=2) ** 2).sum().item()

            x = input_batch.new_zeros(batch_size, dict_size, M, N)
            t = 1
            y = x.clone()
            for i_iter in range(nb_iter):
                x_old = x.clone()
                D_y = nn.functional.conv2d(y, dictionary)
                y = y + nn.functional.conv2d(input_batch - D_y, dictionary.transpose(0, 1).contiguous()) / L
                x = relu(y, lambda_star / L)
                t0 = t
                t = (1. + math.sqrt(1. + 4. * t ** 2)) / 2.
                y = x + ((t0 - 1.) / t) * (x - x_old)

                rec_error = 0.5 * ((nn.functional.conv2d(x, dictionary) - input_batch).norm(p=2, dim=1) ** 2).sum().item()
                sparsity_loss = (lambda_star * x).norm(p=1, dim=1).sum().item()
                loss_curve[i_iter] += rec_error + sparsity_loss
                conv[i_iter] += ((x - x_star).norm(p=2) ** 2).sum().item()

        loss_curve /= (M * N * len(val_loader.dataset))
        conv_rel = conv / x_star_norm

    return loss_curve, conv_rel
//  [python]
def square_root(mat):
    e,v = torch.symeig(mat)
//  [python]
def square_root(mat):
    e,v = torch.linalg.eigh(mat)
//  [python]
def principal_components(coords: torch.Tensor):
        mean = coords.mean(1, keepdim=True)
        x = coords - mean
        n = coords.size()[1]
        c = x.transpose(1, 2) @ x / n
        eig, eiv = torch.symeig(c, eigenvectors=True)
        eivp = eiv.inverse()
        x = (eivp @ x.transpose(1, 2)).transpose(1, 2)
        # x += mean
        return x, eiv, mean
//  [python]
def principal_components(coords: torch.Tensor):
        mean = coords.mean(1, keepdim=True)
        x = coords - mean
        n = coords.size()[1]
        c = x.transpose(1, 2) @ x / n
        eig, eiv = torch.linalg.eigh(c)
        eivp = eiv.inverse()
        x = (eivp @ x.transpose(1, 2)).transpose(1, 2)
        # x += mean
        return x, eiv, mean
//  [python]
def test_SYMEIG_rank_deficient():
    m = 50
    r = 10
    M = torch.rand((m, m), dtype=torch.float64)
    M = M + M.t()
    D, U = torch.symeig(M, eigenvectors=True)
    D[-r:] = 0
    M = U @ torch.diag(D) @ U.t()

    D, U = SYMEIG.apply(M)
    assert torch.norm(M - U @ torch.diag(D) @ U.t()) < D[0] * (M.size()[0] ** 2) * 1e-14

    M.requires_grad_(True)

    def force_sym_eig(M):
        M = 0.5 * (M + M.t())
        D, U = SYMEIG.apply(M)
        return U
    assert torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4)
//  [python]
def test_SYMEIG_rank_deficient():
    m = 50
    r = 10
    M = torch.rand((m, m), dtype=torch.float64)
    M = M + M.t()
    D, U = torch.linalg.eigh(M)
    D[-r:] = 0
    M = U @ torch.diag(D) @ U.t()

    D, U = SYMEIG.apply(M)
    assert torch.norm(M - U @ torch.diag(D) @ U.t()) < D[0] * (M.size()[0] ** 2) * 1e-14

    M.requires_grad_(True)

    def force_sym_eig(M):
        M = 0.5 * (M + M.t())
        D, U = SYMEIG.apply(M)
        return U
    assert torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4)
//  [python]
def test_symeig_removed_error(self, device):
        a = make_tensor(5, 5, device=device, dtype=torch.float32)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            torch.symeig(a)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            a.symeig()
//  [python]
def test_symeig_removed_error(self, device):
        a = make_tensor(5, 5, device=device, dtype=torch.float32)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            torch.linalg.eigh(a)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            torch.linalg.eigh(a)
//  [python]
def symeig(y: torch.Tensor, eigenvectors=True):
    """
    If Y is a squared symmetric matrix, then Y can be decomposed (diagonalized) as Y = SDS^-1
    where S is the matrix composed by the eigenvectors of Y (S = [v_1, v_2, ..., v_n] where v_i are the
    eigenvectors), D is the diagonal matrix constructed from the corresponding eigenvalues, and S^-1 is the
    matrix inverse of S.

    :param y: b x * x n x n. PRE: Each matrix must be symmetric
    :return: eigenvalues (in ascending order): b x * x n
    :return: eigenvectors: b x * x n x n
    """
    return torch.symeig(y, eigenvectors=eigenvectors)
//  [python]
def symeig(y: torch.Tensor, eigenvectors=True):
    """
    If Y is a squared symmetric matrix, then Y can be decomposed (diagonalized) as Y = SDS^-1
    where S is the matrix composed by the eigenvectors of Y (S = [v_1, v_2, ..., v_n] where v_i are the
    eigenvectors), D is the diagonal matrix constructed from the corresponding eigenvalues, and S^-1 is the
    matrix inverse of S.

    :param y: b x * x n x n. PRE: Each matrix must be symmetric
    :return: eigenvalues (in ascending order): b x * x n
    :return: eigenvectors: b x * x n x n
    """
    return torch.linalg.eigh(y, eigenvectors=eigenvectors)
//  [python]
def symeig(M):
    """Symetric eigendecomposition avoiding failure cases by
    adding and removing jitter to the diagonal
    returns eigenvalues (l) and eigenvectors (W)
    """
    # could make double to get more precise computation
    # M = M.double()
    # and then below return L.float(), W.float()
    try:
        L, W = torch.symeig(M, eigenvectors=True)
    except RuntimeError:  # did not converge
        logging.info('SYMEIG: adding jitter, did not converge.')
        # use W L W^T + I = W (L + I) W^T
        M = diagonal_add_scalar(M, value=1.)
        try:
            L, W = torch.symeig(M, eigenvectors=True)
            L -= 1.
        except RuntimeError:
            stats = f'diag: {M.diagonal()}, max: {M.abs().max()}, '
            stats = stats + f'min: {M.abs().min()}, mean: {M.abs().mean()}'
            logging.info(f'SYMEIG: adding jitter failed. Stats: {stats}')
            exit()
    # eigenvalues of symeig at least 0
    L = L.clamp(min=0.0)
    # underflow leads to nans instead of zeros
    L[torch.isnan(L)] = 0.0
    W[torch.isnan(W)] = 0.0
    return L, W
//  [python]
def symeig(M):
    """Symetric eigendecomposition avoiding failure cases by
    adding and removing jitter to the diagonal
    returns eigenvalues (l) and eigenvectors (W)
    """
    # could make double to get more precise computation
    # M = M.double()
    # and then below return L.float(), W.float()
    try:
        L, W = torch.linalg.eigh(M)
    except RuntimeError:  # did not converge
        logging.info('SYMEIG: adding jitter, did not converge.')
        # use W L W^T + I = W (L + I) W^T
        M = diagonal_add_scalar(M, value=1.)
        try:
            L, W = torch.linalg.eigh(M)
            L -= 1.
        except RuntimeError:
            stats = f'diag: {M.diagonal()}, max: {M.abs().max()}, '
            stats = stats + f'min: {M.abs().min()}, mean: {M.abs().mean()}'
            logging.info(f'SYMEIG: adding jitter failed. Stats: {stats}')
            exit()
    # eigenvalues of symeig at least 0
    L = L.clamp(min=0.0)
    # underflow leads to nans instead of zeros
    L[torch.isnan(L)] = 0.0
    W[torch.isnan(W)] = 0.0
    return L, W
//  [python]
def entropy(self, *args):

        for idx, val in enumerate(args):
            if idx == 0:
                k = val.clone()
            else:
                k *= val

        k /= k.trace()
        eigv = th.symeig(k)[0].abs()

        return -(eigv*(eigv.log2())).sum()
//  [python]
def entropy(self, *args):

        for idx, val in enumerate(args):
            if idx == 0:
                k = val.clone()
            else:
                k *= val

        k /= k.trace()
        eigv = th.linalg.eigh(k)[0].abs()

        return -(eigv*(eigv.log2())).sum()
//  [python]
def loss_8a(self, G):
        """
        G: B x 6 x 3N
        """
        Gt = G.transpose(1,2)
        temp = self.eps * self.eye6
        temp = torch.matmul(G, Gt) - temp
        eigval = torch.symeig(temp.cpu(), eigenvectors=True)[0].to(self.device)
        rnev = self.relu(-eigval)
        result = torch.sum(rnev * rnev, 1)
        return result
//  [python]
def loss_8a(self, G):
        """
        G: B x 6 x 3N
        """
        Gt = G.transpose(1,2)
        temp = self.eps * self.eye6
        temp = torch.matmul(G, Gt) - temp
        eigval = torch.linalg.eigh(temp.cpu())[0].to(self.device)
        rnev = self.relu(-eigval)
        result = torch.sum(rnev * rnev, 1)
        return result
//  [python]
def nearestPDHack(c):
    """ this probably kinda works but is really slow"""
    eps = torch.finfo(c.dtype).eps
    k = 1
    while not isPD(c):
        # fix it so we have positive definite matrix
        # could also use the Higham algorithm for more accuracy
        #  N.J. Higham, "Computing a nearest symmetric positive semidefinite
        # https://gist.github.com/fasiha/fdb5cec2054e6f1c6ae35476045a0bbd
        print('covariance matrix not positive definite, attempting recovery')
        e, v = torch.symeig(c, eigenvectors=True)
        bump = eps * k ** 2
        e[e < bump] += bump
        c = torch.matmul(v, torch.matmul(e.diag_embed(), v.t()))
        k += 1

    return c
//  [python]
def nearestPDHack(c):
    """ this probably kinda works but is really slow"""
    eps = torch.finfo(c.dtype).eps
    k = 1
    while not isPD(c):
        # fix it so we have positive definite matrix
        # could also use the Higham algorithm for more accuracy
        #  N.J. Higham, "Computing a nearest symmetric positive semidefinite
        # https://gist.github.com/fasiha/fdb5cec2054e6f1c6ae35476045a0bbd
        print('covariance matrix not positive definite, attempting recovery')
        e, v = torch.linalg.eigh(c)
        bump = eps * k ** 2
        e[e < bump] += bump
        c = torch.matmul(v, torch.matmul(e.diag_embed(), v.t()))
        k += 1

    return c
//  [python]
def pick(self, idx=None):
        # If idx is None, randomly pick
        if idx == None:
            data_idx = random.randint(0, len(self.test_data_ls) - 1)
        else:
            data_idx = idx
        
        G, label = self.test_data_ls[data_idx]
        label = label.to(self.device)
        
        # Input
        X0 = G.ndata['feat'].to(self.device)
        E0 = G.edata['feat'].to(self.device) if 'feat' in G.edata else torch.zeros(0)
        
        # Eigen decomposition
        A = G.adjacency_matrix(transpose=True).to_dense().to(dtype=torch.float64)
        A[A > 0] = 1.0

        D_inv = torch.sum(A, dim=1)
        D_inv[torch.where(D_inv > 0)] = 1 / D_inv[np.where(D_inv > 0)]
        D_inv = torch.diag(D_inv)
        A_n = torch.sqrt(D_inv) @ A @ torch.sqrt(D_inv ** 0.5)

        Lam, U = torch.symeig(A_n, eigenvectors=True)
        U = U.type(torch.DoubleTensor)
        
        return (G, label, X0, E0, U)
//  [python]
def pick(self, idx=None):
        # If idx is None, randomly pick
        if idx == None:
            data_idx = random.randint(0, len(self.test_data_ls) - 1)
        else:
            data_idx = idx
        
        G, label = self.test_data_ls[data_idx]
        label = label.to(self.device)
        
        # Input
        X0 = G.ndata['feat'].to(self.device)
        E0 = G.edata['feat'].to(self.device) if 'feat' in G.edata else torch.zeros(0)
        
        # Eigen decomposition
        A = G.adjacency_matrix(transpose=True).to_dense().to(dtype=torch.float64)
        A[A > 0] = 1.0

        D_inv = torch.sum(A, dim=1)
        D_inv[torch.where(D_inv > 0)] = 1 / D_inv[np.where(D_inv > 0)]
        D_inv = torch.diag(D_inv)
        A_n = torch.sqrt(D_inv) @ A @ torch.sqrt(D_inv ** 0.5)

        Lam, U = torch.linalg.eigh(A_n)
        U = U.type(torch.DoubleTensor)
        
        return (G, label, X0, E0, U)
//  [python]
def dist(self, a: torch.Tensor, b: torch.Tensor, keepdim=True) -> Tuple[torch.Tensor, torch.Tensor]:
        """Compute distance between 2 points on the SPD manifold

        M = A^(-1/2) B A^(-1/2)
        UDU^1 = M
        D: matrix with eigenvalues of m in descending order: d1, d2,..., dn, where d1 > d2 > ... > dn
        log_eigs = log([d1, d2, ..., dn])
        And then on log_eigs different metrics can be computed.

        :param a, b: b x n x n: elements in SPD
        :return: distance: b: torch.Tensor with distances between a and b in SPD
        :return: vvd: b x n: torch.Tensor with vector-valued distance between a and b
        """
        inv_sqrt_a = lalg.sym_inv_sqrtm1(a)
        m = inv_sqrt_a @ b @ inv_sqrt_a
        m = lalg.sym(m)         # impose symmetry due to numeric instabilities

        try:
            eigvals, _ = torch.symeig(m, eigenvectors=True)     # eigvalues are in descending order
        except RuntimeError as e:
            log = get_logging()
            log.info(f"ERROR: torch.symeig in SPD dist did not converge. m = {m}")
            raise e

        log_eigvals = torch.log(eigvals)
        result = self.metric.compute_metric(log_eigvals, keepdim=keepdim)
        return result, log_eigvals
//  [python]
def dist(self, a: torch.Tensor, b: torch.Tensor, keepdim=True) -> Tuple[torch.Tensor, torch.Tensor]:
        """Compute distance between 2 points on the SPD manifold

        M = A^(-1/2) B A^(-1/2)
        UDU^1 = M
        D: matrix with eigenvalues of m in descending order: d1, d2,..., dn, where d1 > d2 > ... > dn
        log_eigs = log([d1, d2, ..., dn])
        And then on log_eigs different metrics can be computed.

        :param a, b: b x n x n: elements in SPD
        :return: distance: b: torch.Tensor with distances between a and b in SPD
        :return: vvd: b x n: torch.Tensor with vector-valued distance between a and b
        """
        inv_sqrt_a = lalg.sym_inv_sqrtm1(a)
        m = inv_sqrt_a @ b @ inv_sqrt_a
        m = lalg.sym(m)         # impose symmetry due to numeric instabilities

        try:
            eigvals, _ = torch.linalg.eigh(m)     # eigvalues are in descending order
        except RuntimeError as e:
            log = get_logging()
            log.info(f"ERROR: torch.linalg.eigh in SPD dist did not converge. m = {m}")
            raise e

        log_eigvals = torch.log(eigvals)
        result = self.metric.compute_metric(log_eigvals, keepdim=keepdim)
        return result, log_eigvals
//  [python]
def compute_condnum(net, inputs, targets, mode, loss_fn, split_data=1):
    net.zero_grad()
    N = inputs.shape[0]
    grads = []

    autograd_hacks.add_hooks(net)
    outputs = net.forward(inputs)
    sum(outputs[torch.arange(N), targets]).backward()
    autograd_hacks.compute_grad1(net, loss_type='sum')

    grads = [
        param.grad1.flatten(start_dim=1) for param in net.parameters()
        if hasattr(param, 'grad1')
    ]
    grads = torch.cat(grads, axis=1)

    ntk = torch.matmul(grads, grads.t())
    eigenvalues, _ = torch.symeig(ntk)  # ascending
    return np.nan_to_num(
        (eigenvalues[-1] / eigenvalues[0]).item(), copy=True, nan=100000.0)
//  [python]
def compute_condnum(net, inputs, targets, mode, loss_fn, split_data=1):
    net.zero_grad()
    N = inputs.shape[0]
    grads = []

    autograd_hacks.add_hooks(net)
    outputs = net.forward(inputs)
    sum(outputs[torch.arange(N), targets]).backward()
    autograd_hacks.compute_grad1(net, loss_type='sum')

    grads = [
        param.grad1.flatten(start_dim=1) for param in net.parameters()
        if hasattr(param, 'grad1')
    ]
    grads = torch.cat(grads, axis=1)

    ntk = torch.matmul(grads, grads.t())
    eigenvalues, _ = torch.linalg.eigh(ntk)  # ascending
    return np.nan_to_num(
        (eigenvalues[-1] / eigenvalues[0]).item(), copy=True, nan=100000.0)
//  [python]
def projected_pca(JtJ, bases=None):
    if bases is not None:
        JtJ = bases.T @ JtJ @ bases
    S, V = torch.symeig(JtJ, eigenvectors=True)
    index = S.abs().argsort(descending=True)
    S, V = S[index], V[:, index]
    if bases is not None:
        V = bases @ V
    return S, V
//  [python]
def projected_pca(JtJ, bases=None):
    if bases is not None:
        JtJ = bases.T @ JtJ @ bases
    S, V = torch.linalg.eigh(JtJ)
    index = S.abs().argsort(descending=True)
    S, V = S[index], V[:, index]
    if bases is not None:
        V = bases @ V
    return S, V
//  [python]
def loss(self, H1, H2):
        r1 = 1e-3
        r2 = 1e-3
        eps = 1e-9

        H1, H2 = H1.t(), H2.t()
        o1 = o2 = H1.size(0)

        m = H1.size(1)

        H1bar = H1 - H1.mean(dim=1).unsqueeze(dim=1)
        H2bar = H2 - H2.mean(dim=1).unsqueeze(dim=1)

        SigmaHat12 = (1.0 / (m - 1)) * torch.matmul(H1bar, H2bar.t())
        SigmaHat11 = (1.0 / (m - 1)) * torch.matmul(H1bar,
                                                    H1bar.t()) + r1 * torch.eye(o1, device=self.device)
        SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar,
                                                    H2bar.t()) + r2 * torch.eye(o2, device=self.device)

        # Calculating the root inverse of covariance matrices by using eigen decomposition
        breakpoint()
        [D1, V1] = torch.symeig(SigmaHat11, eigenvectors=True)
        [D2, V2] = torch.symeig(SigmaHat22, eigenvectors=True)

        # Added to increase stability
        posInd1 = torch.gt(D1, eps).nonzero()[:, 0]
        D1 = D1[posInd1]
        V1 = V1[:, posInd1]
        posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
        D2 = D2[posInd2]
        V2 = V2[:, posInd2]

        SigmaHat11RootInv = torch.matmul(
            torch.matmul(V1, torch.diag(D1 ** -0.5)), V1.t())
        SigmaHat22RootInv = torch.matmul(
            torch.matmul(V2, torch.diag(D2 ** -0.5)), V2.t())

        Tval = torch.matmul(torch.matmul(SigmaHat11RootInv,
                                         SigmaHat12), SigmaHat22RootInv)

        if self.use_all_singular_values:
            # all singular values are used to calculate the correlation
            tmp = torch.matmul(Tval.t(), Tval)
            corr = torch.trace(torch.sqrt(tmp))
            # assert torch.isnan(corr).item() == 0
        else:
            # just the top self.outdim_size singular values are used
            trace_TT = torch.matmul(Tval.t(), Tval)
            trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0])*r1).to(self.device)) # regularization for more stability
            U, V = torch.symeig(trace_TT, eigenvectors=True)
            U = torch.where(U>eps, U, (torch.ones(U.shape)*eps).to(self.device))
            U = U.topk(self.outdim_size)[0]
            corr = torch.sum(torch.sqrt(U))
        return -corr
//  [python]
def loss(self, H1, H2):
        r1 = 1e-3
        r2 = 1e-3
        eps = 1e-9

        H1, H2 = H1.t(), H2.t()
        o1 = o2 = H1.size(0)

        m = H1.size(1)

        H1bar = H1 - H1.mean(dim=1).unsqueeze(dim=1)
        H2bar = H2 - H2.mean(dim=1).unsqueeze(dim=1)

        SigmaHat12 = (1.0 / (m - 1)) * torch.matmul(H1bar, H2bar.t())
        SigmaHat11 = (1.0 / (m - 1)) * torch.matmul(H1bar,
                                                    H1bar.t()) + r1 * torch.eye(o1, device=self.device)
        SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(H2bar,
                                                    H2bar.t()) + r2 * torch.eye(o2, device=self.device)

        # Calculating the root inverse of covariance matrices by using eigen decomposition
        [D1, V1] = torch.linalg.eigh(SigmaHat11)
        [D2, V2] = torch.linalg.eigh(SigmaHat22)

        # Added to increase stability
        posInd1 = torch.gt(D1, eps).nonzero()[:, 0]
        D1 = D1[posInd1]
        V1 = V1[:, posInd1]
        posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
        D2 = D2[posInd2]
        V2 = V2[:, posInd2]

        SigmaHat11RootInv = torch.matmul(
            torch.matmul(V1, torch.diag(D1 ** -0.5)), V1.t())
        SigmaHat22RootInv = torch.matmul(
            torch.matmul(V2, torch.diag(D2 ** -0.5)), V2.t())

        Tval = torch.matmul(torch.matmul(SigmaHat11RootInv,
                                         SigmaHat12), SigmaHat22RootInv)

        if self.use_all_singular_values:
            # all singular values are used to calculate the correlation
            tmp = torch.matmul(Tval.t(), Tval)
            corr = torch.trace(torch.sqrt(tmp))
            # assert torch.isnan(corr).item() == 0
        else:
            # just the top self.outdim_size singular values are used
            trace_TT = torch.matmul(Tval.t(), Tval)
            trace_TT = torch.add(trace_TT, (torch.eye(trace_TT.shape[0])*r1).to(self.device)) # regularization for more stability
            U, V = torch.linalg.eigh(trace_TT)
            U = torch.where(U>eps, U, (torch.ones(U.shape)*eps).to(self.device))
            U = U.topk(self.outdim_size)[0]
            corr = torch.sum(torch.sqrt(U))
        return -corr
//  [python]
def forward(ctx, input, eps=1e-2):
        use_cuda = input.is_cuda
        #if input.size(0) < 300:
        #    input = input.cpu()
        input = input.cpu()
        #print(torch.isnan(input).any())
        e, v = torch.symeig(input, eigenvectors=True)
        if use_cuda:
            e = e.cuda()
            v = v.cuda()
        e.clamp_(min=0)
        e_sqrt = e.sqrt_().add_(eps)
        ctx.save_for_backward(e_sqrt, v)
        e_rsqrt = e_sqrt.reciprocal()

        output = v.mm(torch.diag(e_rsqrt).mm(v.t()))
        return output
//  [python]
def forward(ctx, input, eps=1e-2):
        use_cuda = input.is_cuda
        #if input.size(0) < 300:
        #    input = input.cpu()
        input = input.cpu()
        #print(torch.isnan(input).any())
        e, v = torch.linalg.eigh(input)
        if use_cuda:
            e = e.cuda()
            v = v.cuda()
        e.clamp_(min=0)
        e_sqrt = e.sqrt_().add_(eps)
        ctx.save_for_backward(e_sqrt, v)
        e_rsqrt = e_sqrt.reciprocal()

        output = v.mm(torch.diag(e_rsqrt).mm(v.t()))
        return output
//  [python]
def get_eigenvalues(
        factors: List[Tensor],
        verbose: bool = False) -> Tensor:
    """Computes the eigenvalues of KFAC, EFB or diagonal factors.

    Args:
        factors: A list of KFAC, EFB or diagonal factors.
        verbose: Prints out progress if True.

    Returns:
        The eigenvalues of all KFAC, EFB or diagonal factors.
    """
    eigenvalues = Tensor()
    factors = tqdm(factors, disable=not verbose)
    for layer, factor in enumerate(factors):
        factors.set_description(desc=f"Layer [{layer + 1}/{len(factors)}]")
        if len(factor) == 2:
            xxt_eigvals = torch.symeig(factor[0])[0]
            ggt_eigvals = torch.symeig(factor[1])[0]
            eigenvalues = torch.cat([eigenvalues, torch.ger(xxt_eigvals, ggt_eigvals).contiguous().view(-1)])
        else:
            eigenvalues = torch.cat([eigenvalues, factor.contiguous().view(-1)])
    return eigenvalues
//  [python]
def get_eigenvalues(
        factors: List[Tensor],
        verbose: bool = False) -> Tensor:
    """Computes the eigenvalues of KFAC, EFB or diagonal factors.

    Args:
        factors: A list of KFAC, EFB or diagonal factors.
        verbose: Prints out progress if True.

    Returns:
        The eigenvalues of all KFAC, EFB or diagonal factors.
    """
    eigenvalues = Tensor()
    factors = tqdm(factors, disable=not verbose)
    for layer, factor in enumerate(factors):
        factors.set_description(desc=f"Layer [{layer + 1}/{len(factors)}]")
        if len(factor) == 2:
            xxt_eigvals = torch.linalg.eigh(factor[0])[0]
            ggt_eigvals = torch.linalg.eigh(factor[1])[0]
            eigenvalues = torch.cat([eigenvalues, torch.ger(xxt_eigvals, ggt_eigvals).contiguous().view(-1)])
        else:
            eigenvalues = torch.cat([eigenvalues, factor.contiguous().view(-1)])
    return eigenvalues
//  [python]
def train(cls, data, speaker_labels, device): ## Expects centered data
        print('Training LDA...')
        data = data.to(device)
        data, class_boundaries = _rearrange_data(data, speaker_labels)
        print('Computing within class covariance...')
        Sw = _compute_within_cov(data, class_boundaries)
        print('Computing between class covariance...')
        Sb = _compute_between_cov(data, class_boundaries)
        print('Simultaneous diagonalization...')
        l, U = torch.symeig(Sw, eigenvectors=True)
        W = torch.rsqrt(l) * U
        Sbb = torch.chain_matmul(W.t(), Sb, W)
        l, U = torch.symeig(Sbb, eigenvectors=True)
        W = torch.matmul(W, U) 
        print('LDA trained!...')
        return Lda(W)
//  [python]
def train(cls, data, speaker_labels, device): ## Expects centered data
        print('Training LDA...')
        data = data.to(device)
        data, class_boundaries = _rearrange_data(data, speaker_labels)
        print('Computing within class covariance...')
        Sw = _compute_within_cov(data, class_boundaries)
        print('Computing between class covariance...')
        Sb = _compute_between_cov(data, class_boundaries)
        print('Simultaneous diagonalization...')
        l, U = torch.linalg.eigh(Sw)
        W = torch.rsqrt(l) * U
        Sbb = torch.chain_matmul(W.t(), Sb, W)
        l, U = torch.linalg.eigh(Sbb)
        W = torch.matmul(W, U) 
        print('LDA trained!...')
        return Lda(W)
//  [python]
def get_sqrt_inverse(matrix, bias=1.0e-5):
    '''Compute the sqrt-inverse of the supplied symmetric/real matrix.

    We need this step for whitening and TICA.

    Arguments:
        matrix (Tensor): contains the matrix you want to transform
        bias (float): assures numerical stability
    '''
    e, v = _torch.symeig(matrix, eigenvectors=True)
    d = _torch.diag(1.0 / _torch.sqrt(_torch.abs(e) + bias))
    return _torch.mm(_torch.mm(v, d), v.t())
//  [python]
def get_sqrt_inverse(matrix, bias=1.0e-5):
    '''Compute the sqrt-inverse of the supplied symmetric/real matrix.

    We need this step for whitening and TICA.

    Arguments:
        matrix (Tensor): contains the matrix you want to transform
        bias (float): assures numerical stability
    '''
    e, v = _torch.linalg.eigh(matrix)
    d = _torch.diag(1.0 / _torch.sqrt(_torch.abs(e) + bias))
    return _torch.mm(_torch.mm(v, d), v.t())
//  [python]
def forward(self, x):
        N = x.shape[0]
        # if self.x_last_batch is None:
        #     self.x_last_batch = torch.randn_like(x)
        # x, self.x_last_batch = torch.cat([x, self.x_last_batch]), x.detach()
        mean = x.mean(dim=0)
        if self.running_mean is None:
            self.running_mean = mean
        else:
            mean = self.running_mean = (1. - self.momentum) * self.running_mean.detach() + self.momentum * mean

        x = x - mean
        # import pdb
        # pdb.set_trace()
        cov = x.t().matmul(x) / (x.size(0) - 1)
        if self.running_covariance is None:
            self.running_covariance = cov
        else:
            cov = self.running_covariance = (1 - self.momentum) * self.running_covariance.detach() + self.momentum * cov
        eigenvalues, eigenvectors = torch.symeig(cov.cpu(), eigenvectors=True, upper=True)
        S, U = eigenvalues.to(x.device), eigenvectors.to(x.device)
        self.eig = eigenvalues.min()
        whitening_transform = U.matmul(S.rsqrt().diag()).matmul(U.t())

        return x.matmul(whitening_transform)
//  [python]
def forward(self, x):
        N = x.shape[0]
        # if self.x_last_batch is None:
        #     self.x_last_batch = torch.randn_like(x)
        # x, self.x_last_batch = torch.cat([x, self.x_last_batch]), x.detach()
        mean = x.mean(dim=0)
        if self.running_mean is None:
            self.running_mean = mean
        else:
            mean = self.running_mean = (1. - self.momentum) * self.running_mean.detach() + self.momentum * mean

        x = x - mean
        # import pdb
        # pdb.set_trace()
        cov = x.t().matmul(x) / (x.size(0) - 1)
        if self.running_covariance is None:
            self.running_covariance = cov
        else:
            cov = self.running_covariance = (1 - self.momentum) * self.running_covariance.detach() + self.momentum * cov
        eigenvalues, eigenvectors = torch.linalg.eigh(cov.cpu())
        S, U = eigenvalues.to(x.device), eigenvectors.to(x.device)
        self.eig = eigenvalues.min()
        whitening_transform = U.matmul(S.rsqrt().diag()).matmul(U.t())

        return x.matmul(whitening_transform)
//  [python]
def _t_make_grad_eig_x(self, t_dm):
        # v stands for eigen values
        shell_pdm = [torch.einsum('rap,rs,saq->apq', po, t_dm, po).requires_grad_(True)
                        for po in self._t_ovlp_shells]
        calc_eig = lambda dm: torch.symeig(dm, True)[0]
        shell_gvdm = [get_batch_jacobian(calc_eig, dm, dm.shape[-1]) 
                        for dm in shell_pdm]
        shell_gdmx = self._t_make_grad_pdm_x(t_dm)
        shell_gvx = [torch.einsum("bxapq,avpq->bxav", gdmx, gvdm) 
                        for gdmx, gvdm in zip(shell_gdmx, shell_gvdm)]
        return torch.cat(shell_gvx, dim=-1)
//  [python]
def _t_make_grad_eig_x(self, t_dm):
        # v stands for eigen values
        shell_pdm = [torch.einsum('rap,rs,saq->apq', po, t_dm, po).requires_grad_(True)
                        for po in self._t_ovlp_shells]
        calc_eig = lambda dm: torch.linalg.eigh(dm, UPLO='U')[0]
        shell_gvdm = [get_batch_jacobian(calc_eig, dm, dm.shape[-1]) 
                        for dm in shell_pdm]
        shell_gdmx = self._t_make_grad_pdm_x(t_dm)
        shell_gvx = [torch.einsum("bxapq,avpq->bxav", gdmx, gvdm) 
                        for gdmx, gvdm in zip(shell_gdmx, shell_gvdm)]
        return torch.cat(shell_gvx, dim=-1)
//  [python]
def whiten_col(tx,eps=1e-4):
    shape = tx.shape
    tx = tx.squeeze()
    D = torch.mm(tx, tx.t()) / len(tx)
    diag, v = torch.symeig(D, eigenvectors=True)
    diag[diag < eps] = 1
    diag = diag ** 0.5
    diag = 1 / diag
    S = torch.diag(diag)
    out = v @ S @ v.t() @ tx
    out = out.view(shape)
    return out
//  [python]
def whiten_col(tx,eps=1e-4):
    shape = tx.shape
    tx = tx.squeeze()
    D = torch.mm(tx, tx.t()) / len(tx)
    diag, v = torch.linalg.eigh(D)
    diag[diag < eps] = 1
    diag = diag ** 0.5
    diag = 1 / diag
    S = torch.diag(diag)
    out = v @ S @ v.t() @ tx
    out = out.view(shape)
    return out
//  [python]
def _t_get_grad_dms(mf, dm=None):
    # calculate \partial E / \partial (D^I_rl)_mm' by shells
    if dm is None:
        dm = mf.make_rdm1()
    t_dm = torch.from_numpy(dm).double().to(mf.device)
    proj_dms = [torch.einsum('rap,rs,saq->apq', po, t_dm, po).requires_grad_(True)
                    for po in mf._t_ovlp_shells]
    if mf.net is None:
        return [torch.zeros_like(pdm) for pdm in proj_dms]
    proj_eigs = [torch.symeig(dm, eigenvectors=True)[0]
                    for dm in proj_dms]
    ceig = torch.cat(proj_eigs, dim=-1).unsqueeze(0) # 1 x natoms x nproj
    ec = mf.net(ceig)
    grad_dms = torch.autograd.grad(ec, proj_dms)
    return grad_dms
//  [python]
def _t_get_grad_dms(mf, dm=None):
    # calculate \partial E / \partial (D^I_rl)_mm' by shells
    if dm is None:
        dm = mf.make_rdm1()
    t_dm = torch.from_numpy(dm).double().to(mf.device)
    proj_dms = [torch.einsum('rap,rs,saq->apq', po, t_dm, po).requires_grad_(True)
                    for po in mf._t_ovlp_shells]
    if mf.net is None:
        return [torch.zeros_like(pdm) for pdm in proj_dms]
    proj_eigs = [torch.linalg.eigh(dm)[0]
                    for dm in proj_dms]
    ceig = torch.cat(proj_eigs, dim=-1).unsqueeze(0) # 1 x natoms x nproj
    ec = mf.net(ceig)
    grad_dms = torch.autograd.grad(ec, proj_dms)
    return grad_dms
//  [python]
def train(self, x):

        print('calculating PCA parameters...')

        x = x.t()
        nPoints = x.size(1)
        nDims = x.size(0)

        mu = x.mean(1).unsqueeze(1)
        x = x - mu

        if (nDims<=nPoints):
            doDual = False
            x2 = torch.matmul(x, x.t()) / (nPoints - 1)
        else:
            doDual = True
            x2 = torch.matmul(x.t(), x) / (nPoints - 1)

        L, U = torch.symeig(x2, eigenvectors=True)
        if (self.pca_n_components < x2.size(0)):
            k_indices = torch.argsort(L, descending=True)[:self.pca_n_components]
            L = torch.index_select(L, 0, k_indices)
            U = torch.index_select(U, 1, k_indices)

        lams = L
        lams[lams<1e-9] = 1e-9

        if (doDual):
            U = torch.matmul(x, torch.matmul(U, torch.diag(1./torch.sqrt(lams))/np.sqrt(nPoints-1)))

        Utmu= torch.matmul(U.t(), mu)

        U, lams, mu, Utmu = U.numpy(), lams.numpy(), mu.numpy(), Utmu.numpy()

        print('================= PCA RESULT ==================')
        print('U: {}'.format(U.shape))
        print('lams: {}'.format(lams.shape))
        print('mu: {}'.format(mu.shape))
        print('Utmu: {}'.format(Utmu.shape))
        print('===============================================')

        filename = os.path.join(self.pca_parameters_path)
        h5file = h5py.File(filename, 'w')
        h5file.create_dataset('U', data=U)
        h5file.create_dataset('lams', data=lams)
        h5file.create_dataset('mu', data=mu)
        h5file.create_dataset('Utmu', data=Utmu)
        h5file.close()
//  [python]
def train(self, x):

        print('calculating PCA parameters...')

        x = x.t()
        nPoints = x.size(1)
        nDims = x.size(0)

        mu = x.mean(1).unsqueeze(1)
        x = x - mu

        if (nDims<=nPoints):
            doDual = False
            x2 = torch.matmul(x, x.t()) / (nPoints - 1)
        else:
            doDual = True
            x2 = torch.matmul(x.t(), x) / (nPoints - 1)

        L, U = torch.linalg.eigh(x2)
        if (self.pca_n_components < x2.size(0)):
            k_indices = torch.argsort(L, descending=True)[:self.pca_n_components]
            L = torch.index_select(L, 0, k_indices)
            U = torch.index_select(U, 1, k_indices)

        lams = L
        lams[lams<1e-9] = 1e-9

        if (doDual):
            U = torch.matmul(x, torch.matmul(U, torch.diag(1./torch.sqrt(lams))/np.sqrt(nPoints-1)))

        Utmu= torch.matmul(U.t(), mu)

        U, lams, mu, Utmu = U.numpy(), lams.numpy(), mu.numpy(), Utmu.numpy()

        print('================= PCA RESULT ==================')
        print('U: {}'.format(U.shape))
        print('lams: {}'.format(lams.shape))
        print('mu: {}'.format(mu.shape))
        print('Utmu: {}'.format(Utmu.shape))
        print('===============================================')

        filename = os.path.join(self.pca_parameters_path)
        h5file = h5py.File(filename, 'w')
        h5file.create_dataset('U', data=U)
        h5file.create_dataset('lams', data=lams)
        h5file.create_dataset('mu', data=mu)
        h5file.create_dataset('Utmu', data=Utmu)
        h5file.close()
//  [python]
def _symeig(self, eigenvectors: bool = False) -> Tuple[Tensor, Optional["LazyTensor"]]:
        """Method that allows implementing special-cased symeig computation. Should not be called directly"""
        from lazy.non_lazy_tensor import NonLazyTensor

        dtype = self.dtype  # perform decomposition in double precision for numerical stability
        # TODO: Use fp64 registry once #1213 is addressed
        evals, evecs = torch.symeig(self.evaluate().to(dtype=torch.double), eigenvectors=eigenvectors)
        # chop any negative eigenvalues. TODO: warn if evals are significantly negative
        evals = evals.clamp_min(0.0).to(dtype=dtype)
        if eigenvectors:
            evecs = NonLazyTensor(evecs.to(dtype=dtype))
        else:
            evecs = None
        return evals, evecs