jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_eigenvectors(factors: Dict[Module, Tensor]) -> Dict[Module, Tensor]:
    """Computes the eigenvectors of KFAC factors.

    Args:
        factors: A dict mapping layers to lists of first and second KFAC factors.

    Returns:
        A dict mapping layers to lists containing the first and second KFAC factors eigenvectors.
    """
    eigenvectors = dict()
    for layer, (xxt, ggt) in factors.items():
        sym_xxt, sym_ggt = xxt + xxt.t(), ggt + ggt.t()
        _, xxt_eigvecs = torch.symeig(sym_xxt, eigenvectors=True)
        _, ggt_eigvecs = torch.symeig(sym_ggt, eigenvectors=True)
        eigenvectors[layer] = (xxt_eigvecs, ggt_eigvecs)
    return eigenvectors
//  [python]
def get_eigenvectors(factors: Dict[Module, Tensor]) -> Dict[Module, Tensor]:
    """Computes the eigenvectors of KFAC factors.

    Args:
        factors: A dict mapping layers to lists of first and second KFAC factors.

    Returns:
        A dict mapping layers to lists containing the first and second KFAC factors eigenvectors.
    """
    eigenvectors = dict()
    for layer, (xxt, ggt) in factors.items():
        sym_xxt, sym_ggt = xxt + xxt.t(), ggt + ggt.t()
        _, xxt_eigvecs = torch.linalg.eigh(sym_xxt)
        _, ggt_eigvecs = torch.linalg.eigh(sym_ggt)
        eigenvectors[layer] = (xxt_eigvecs, ggt_eigvecs)
    return eigenvectors
//  [python]
def set_diag_e(e, E, E0, C, eps=1e-6):
#    ipdb.set_trace()
    if E is not None:
        e = E
        print('e = {:.2e}, E = {:.2e}'.format(e, E))
    elif tc.is_tensor(E0):
        if C.size(1) == 1:
            e0, _ = tc.symeig((C.t()*(eps/E0)).mm(C), eigenvectors=True)
#            e0 = (C.t()*(E0)).mm(C).squeeze(0)
#            e0[e0<1e-8] = 1e-8
#            e0 = eps/e0
        else:
            e0, _ = tc.symeig((C.t()*(eps/E0)).mm(C), eigenvectors=True)
#            e0[e0<1e-8] = 1e-8
#            e0 = eps/e0
#            e0 = flip(e0)
        if e is None:
            e = e0
        else:
            e = tc.cat((e, e0))
        print('e0 = [{:.2e}, {:.2e}], E0 = [{:.2e}, {:.2e}]'.format(
              e0[0], e0[-1], E0[0], E0[-1]))
    else:
        e = eps
    return e
//  [python]
def set_diag_e(e, E, E0, C, eps=1e-6):
    if E is not None:
        e = E
        print('e = {:.2e}, E = {:.2e}'.format(e, E))
    elif tc.is_tensor(E0):
        if C.size(1) == 1:
            e0, _ = tc.linalg.eigh((C.t()*(eps/E0)).mm(C))
        else:
            e0, _ = tc.linalg.eigh((C.t()*(eps/E0)).mm(C))
        if e is None:
            e = e0
        else:
            e = tc.cat((e, e0))
        print('e0 = [{:.2e}, {:.2e}], E0 = [{:.2e}, {:.2e}]'.format(
              e0[0], e0[-1], E0[0], E0[-1]))
    else:
        e = eps
    return e
//  [python]
def get_diag_e(E0, C, eps=1e-6):
    K, KA = C.size()
#    print('K={}, KA={}, E.size={}'.format(K, KA, E0.size(0)))
#    ipdb.set_trace()
    if E0 is None:
        e0 = e1 = eps
    elif tc.is_tensor(E0):
        if K == KA:
            e0 = eps/E0
            e1 = eps/(1.0-E0)
        else:
            if KA == 1:
                e0, _ = tc.symeig((C.t()*(eps/E0)).mm(C), eigenvectors=True)
                e1, _ = tc.symeig((C.t()*(eps/(1.0-E0))).mm(C), eigenvectors=True)
#                e0 = (C.t()*(E0)).mm(C).squeeze(0)
#                e0[e0<1e-8] = 1e-8
#                e0 = eps/e0
#                e1 = (C.t()*(1.0-E0)).mm(C).squeeze(0)
#                e1[e1<1e-8] = 1e-8
#                e1 = eps/e1
            else:
                e0, _ = tc.symeig((C.t()*(eps/E0)).mm(C), eigenvectors=True)
                e1, _ = tc.symeig((C.t()*(eps/(1.0-E0))).mm(C), eigenvectors=True)
#                e0, _ = tc.symeig((C.t()*(E0)).mm(C), eigenvectors=True)
#                e1, _ = tc.symeig((C.t()*((1.0-E0))).mm(C), eigenvectors=True)
#                e0[e0<1e-8] = 1e-8
#                e1[e1<1e-8] = 1e-8
#                e0 = eps/e0
#                e1 = eps/e1
                e1 = flip(e1)
    else:
        e0 = e1 = E0
    return e0, e1
//  [python]
def get_diag_e(E0, C, eps=1e-6):
    K, KA = C.size()
    if E0 is None:
        e0 = e1 = eps
    elif tc.is_tensor(E0):
        if K == KA:
            e0 = eps/E0
            e1 = eps/(1.0-E0)
        else:
            if KA == 1:
                e0, _ = tc.linalg.eigh((C.t()*(eps/E0)).mm(C))
                e1, _ = tc.linalg.eigh((C.t()*(eps/(1.0-E0))).mm(C))
            else:
                e0, _ = tc.linalg.eigh((C.t()*(eps/E0)).mm(C))
                e1, _ = tc.linalg.eigh((C.t()*(eps/(1.0-E0))).mm(C))
                e1 = flip(e1)
    else:
        e0 = e1 = E0
    return e0, e1
//  [python]
def test_symeig_no_eigenvectors(self):
        A = torch.tensor([[1., 2.], [2., 4.]], dtype=torch.float32, requires_grad=True)
        w, v = torch.symeig(A, eigenvectors=False)
        with self.assertRaisesRegex(RuntimeError, 'is not differentiable'):
            torch.autograd.backward([w, v], [torch.ones_like(w), torch.ones_like(v)])
//  [python]
def test_symeig_no_eigenvectors(self):
        A = torch.tensor([[1., 2.], [2., 4.]], dtype=torch.float32, requires_grad=True)
        w, v = torch.linalg.eigh(A, eigenvectors=False)
        with self.assertRaisesRegex(RuntimeError, 'is not differentiable'):
            torch.autograd.backward([w, v], [torch.ones_like(w), torch.ones_like(v)])
//  [python]
def cal_Z_axis(local_cor, local_weight=None, ref_point=None):
    device = local_cor.device
    B, N, _ = local_cor.shape
    cov_matrix = torch.matmul(local_cor.transpose(-1, -2), local_cor) if local_weight is None \
        else Variable(torch.matmul(local_cor.transpose(-1, -2), local_cor * local_weight), requires_grad=True)
    Z_axis = torch.symeig(cov_matrix, eigenvectors=True)[1][:, :, 0]
    mask = (torch.sum(-Z_axis * ref_point, dim=1) < 0).float().unsqueeze(1)
    Z_axis = Z_axis * (1 - mask) - Z_axis * mask

    return Z_axis
//  [python]
def cal_Z_axis(local_cor, local_weight=None, ref_point=None):
    device = local_cor.device
    B, N, _ = local_cor.shape
    cov_matrix = torch.matmul(local_cor.transpose(-1, -2), local_cor) if local_weight is None \
        else Variable(torch.matmul(local_cor.transpose(-1, -2), local_cor * local_weight), requires_grad=True)
    Z_axis = torch.linalg.eigh(cov_matrix)[1][:, :, 0]
    mask = (torch.sum(-Z_axis * ref_point, dim=1) < 0).float().unsqueeze(1)
    Z_axis = Z_axis * (1 - mask) - Z_axis * mask

    return Z_axis
//  [python]
def symsqrt(a, cond=None, return_rank=False, dtype=torch.float32):
    """Symmetric square root of a positive semi-definite matrix.
    See https://github.com/pytorch/pytorch/issues/25481"""

    s, u = torch.symeig(a, eigenvectors=True)
    cond_dict = {torch.float32: 1e3 * 1.1920929e-07, torch.float64: 1E6 * 2.220446049250313e-16}

    if cond in [None, -1]:
        cond = cond_dict[dtype]

    above_cutoff = (abs(s) > cond * torch.max(abs(s)))

    psigma_diag = torch.sqrt(s[above_cutoff])
    u = u[:, above_cutoff]

    B = u @ torch.diag(psigma_diag) @ u.t()
    if return_rank:
        return B, len(psigma_diag)
    else:
        return B
//  [python]
def symsqrt(a, cond=None, return_rank=False, dtype=torch.float32):
    """Symmetric square root of a positive semi-definite matrix.
    See https://github.com/pytorch/pytorch/issues/25481"""

    s, u = torch.linalg.eigh(a)
    cond_dict = {torch.float32: 1e3 * 1.1920929e-07, torch.float64: 1E6 * 2.220446049250313e-16}

    if cond in [None, -1]:
        cond = cond_dict[dtype]

    above_cutoff = (abs(s) > cond * torch.max(abs(s)))

    psigma_diag = torch.sqrt(s[above_cutoff])
    u = u[:, above_cutoff]

    B = u @ torch.diag(psigma_diag) @ u.t()
    if return_rank:
        return B, len(psigma_diag)
    else:
        return B
//  [python]
def get_ntk_n(networks,
              recalbn=0,
              train_mode=False,
              num_batch=None,
              batch_size=None,
              image_size=None,
              device=None):
    if device is None:
        device = torch.device('cpu')

    ntks = []
    for network in networks:
        if train_mode:
            network.train()
        else:
            network.eval()

    grads = [[] for _ in range(len(networks))]

    for i in range(num_batch):
        inputs = torch.randn((batch_size, 3, image_size, image_size),
                             device=device)
        for net_idx, network in enumerate(networks):
            network.zero_grad()
            inputs_ = inputs.clone()

            logit = network(inputs_)
            if isinstance(logit, tuple):
                logit = logit[1]
            for _idx in range(len(inputs_)):
                logit[_idx:_idx + 1].backward(
                    torch.ones_like(logit[_idx:_idx + 1]), retain_graph=True)
                grad = []
                for name, W in network.named_parameters():
                    if 'weight' in name and W.grad is not None:
                        grad.append(W.grad.view(-1).detach())
                grads[net_idx].append(torch.cat(grad, -1))
                network.zero_grad()
                # if gpu is not None:
                #     torch.cuda.empty_cache()

    ######
    grads = [torch.stack(_grads, 0) for _grads in grads]
    ntks = [torch.einsum('nc,mc->nm', [_grads, _grads]) for _grads in grads]
    conds = []
    for ntk in ntks:
        eigenvalues, _ = torch.symeig(ntk)  # ascending
        conds.append(
            np.nan_to_num(
                (eigenvalues[-1] / eigenvalues[0]).item(), copy=True))
    return conds
//  [python]
def get_ntk_n(networks,
              recalbn=0,
              train_mode=False,
              num_batch=None,
              batch_size=None,
              image_size=None,
              device=None):
    if device is None:
        device = torch.device('cpu')

    ntks = []
    for network in networks:
        if train_mode:
            network.train()
        else:
            network.eval()

    grads = [[] for _ in range(len(networks))]

    for i in range(num_batch):
        inputs = torch.randn((batch_size, 3, image_size, image_size),
                             device=device)
        for net_idx, network in enumerate(networks):
            network.zero_grad()
            inputs_ = inputs.clone()

            logit = network(inputs_)
            if isinstance(logit, tuple):
                logit = logit[1]
            for _idx in range(len(inputs_)):
                logit[_idx:_idx + 1].backward(
                    torch.ones_like(logit[_idx:_idx + 1]), retain_graph=True)
                grad = []
                for name, W in network.named_parameters():
                    if 'weight' in name and W.grad is not None:
                        grad.append(W.grad.view(-1).detach())
                grads[net_idx].append(torch.cat(grad, -1))
                network.zero_grad()
                # if gpu is not None:
                #     torch.cuda.empty_cache()

    ######
    grads = [torch.stack(_grads, 0) for _grads in grads]
    ntks = [torch.einsum('nc,mc->nm', [_grads, _grads]) for _grads in grads]
    conds = []
    for ntk in ntks:
        eigenvalues, _ = torch.linalg.eigh(ntk)  # ascending
        conds.append(
            np.nan_to_num(
                (eigenvalues[-1] / eigenvalues[0]).item(), copy=True))
    return conds
//  [python]
def update_bays_reg_BDQN(self, log):
        with torch.no_grad():    
            log.debug("######## Start updating bayesian linear layer ########")
            # Update action posterior with formulas: \beta | z,y ~ N(mu_q, cov_q)
            z = self.latent_z.double()
            y = self.train_y.squeeze().cuda()
            s = torch.matmul(z.T, z)
            A = s/self.sigma_n + 1/self.sigma*self.eye
            B = torch.matmul(z.T, y.double())/self.sigma_n
            # A_eig_val, A_eig_vec = torch.symeig(A, eigenvectors=True)
            inv = torch.inverse(A.double())
            self.mu_w[0] = torch.matmul(inv, B).squeeze()
            temp_cov = self.sigma*inv        
            eig_val, eig_vec = torch.symeig(temp_cov, eigenvectors=True)
            # log.debug("After inverse. eigenvalue (pd): {}".format(eig_val[:20]) )
            # log.debug("After inverse. eigenvalue (pd): {}".format(eig_val[-20:]) )
            if torch.any(eig_val < 0):
                self.cov_w[0] = torch.matmul(torch.matmul(eig_vec, torch.diag(torch.abs(eig_val))), torch.t(eig_vec))
            else:
                self.cov_w[0] = temp_cov
//  [python]
def update_bays_reg_BDQN(self, log):
        with torch.no_grad():    
            log.debug("######## Start updating bayesian linear layer ########")
            # Update action posterior with formulas: \beta | z,y ~ N(mu_q, cov_q)
            z = self.latent_z.double()
            y = self.train_y.squeeze().cuda()
            s = torch.matmul(z.T, z)
            A = s/self.sigma_n + 1/self.sigma*self.eye
            B = torch.matmul(z.T, y.double())/self.sigma_n
            # A_eig_val, A_eig_vec = torch.symeig(A, eigenvectors=True)
            inv = torch.inverse(A.double())
            self.mu_w[0] = torch.matmul(inv, B).squeeze()
            temp_cov = self.sigma*inv        
            eig_val, eig_vec = torch.linalg.eigh(temp_cov)
            # log.debug("After inverse. eigenvalue (pd): {}".format(eig_val[:20]) )
            # log.debug("After inverse. eigenvalue (pd): {}".format(eig_val[-20:]) )
            if torch.any(eig_val < 0):
                self.cov_w[0] = torch.matmul(torch.matmul(eig_vec, torch.diag(torch.abs(eig_val))), torch.t(eig_vec))
            else:
                self.cov_w[0] = temp_cov
//  [python]
def get_Laplace_from_pc(ori_pc):
    """
    ori_pc:(B, 3, N)
    """
    #print("shape of ori pc:",ori_pc.shape)
    pc = ori_pc.detach().clone()
    with torch.no_grad():
        # pc = pc.to('cpu').to(torch.double)
        idx = knn(pc, 30)
        pc = pc.transpose(2, 1).contiguous()  #(B, N, 3)
        point_mat = pc.unsqueeze(2) - pc.unsqueeze(1)  #(B, N, N, 3)
        A = torch.exp(-torch.sum(point_mat.square(), dim=3))  #(B, N, N)
        mask = torch.zeros_like(A)
        mask.scatter_(2, idx, 1)
        mask = mask + mask.transpose(2, 1)
        mask[mask>1] = 1
        
        A = A * mask
        D = torch.diag_embed(torch.sum(A, dim=2))
        L = D - A
        e, v = torch.symeig(L, eigenvectors=True)
    return e.to(ori_pc), v.to(ori_pc)
//  [python]
def get_Laplace_from_pc(ori_pc):
    """
    ori_pc:(B, 3, N)
    """
    #print("shape of ori pc:",ori_pc.shape)
    pc = ori_pc.detach().clone()
    with torch.no_grad():
        # pc = pc.to('cpu').to(torch.double)
        idx = knn(pc, 30)
        pc = pc.transpose(2, 1).contiguous()  #(B, N, 3)
        point_mat = pc.unsqueeze(2) - pc.unsqueeze(1)  #(B, N, N, 3)
        A = torch.exp(-torch.sum(point_mat.square(), dim=3))  #(B, N, N)
        mask = torch.zeros_like(A)
        mask.scatter_(2, idx, 1)
        mask = mask + mask.transpose(2, 1)
        mask[mask>1] = 1
        
        A = A * mask
        D = torch.diag_embed(torch.sum(A, dim=2))
        L = D - A
        e, v = torch.linalg.eigh(L)
    return e.to(ori_pc), v.to(ori_pc)
//  [python]
def main():
    n = 10000
    a = sym(torch.rand(n, 2, 2))

    start = time.time()
    _ = torch.symeig(a)
    end = time.time()
    print('Time cpu: ', end - start)

    a = a.cuda()
    start = time.time()
    _ = torch.symeig(a)
    end = time.time()
    print('Time gpu: ', end - start)
//  [python]
def main():
    n = 10000
    a = sym(torch.rand(n, 2, 2))

    start = time.time()
    _ = torch.linalg.eigh(a)
    end = time.time()
    print('Time cpu: ', end - start)

    a = a.cuda()
    start = time.time()
    _ = torch.linalg.eigh(a)
    end = time.time()
    print('Time gpu: ', end - start)
//  [python]
def leading_eigenvector(mat: Tensor, method: str = "power", num_iterations: int = 10) -> Tensor:
    """Compute the leading eigenvector of the matrix.

    Source: https://github.com/XuyangBai/PointDSC/blob/master/models/PointDSC.py

    Args:
        mat (Tensor): the input matrix in the shape of (*, M, M).
        method (str): the method to use: "power", "torch". Default: "power".
        num_iterations (int): the number of iterations for "power" method. Default: 10.

    Returns:
        A Tensor of the leading eigenvector.
    """
    assert method in ["power", "torch"], f"Unsupported method: {method}."

    if method == "power":
        # power iteration algorithm
        leading_eig = torch.ones_like(mat[..., 0:1])  # (*, M, 1)
        leading_eig_last = leading_eig
        for i in range(num_iterations):
            leading_eig = torch.matmul(mat, leading_eig)
            leading_eig = F.normalize(leading_eig, p=2, dim=-2)
            if torch.allclose(leading_eig, leading_eig_last):
                break
            leading_eig_last = leading_eig
        leading_eig = leading_eig.squeeze(-1)
    else:
        _, eigenvectors = torch.symeig(mat, eigenvectors=True)
        leading_eig = eigenvectors[:, :, -1]

    return leading_eig
//  [python]
def leading_eigenvector(mat: Tensor, method: str = "power", num_iterations: int = 10) -> Tensor:
    """Compute the leading eigenvector of the matrix.

    Source: https://github.com/XuyangBai/PointDSC/blob/master/models/PointDSC.py

    Args:
        mat (Tensor): the input matrix in the shape of (*, M, M).
        method (str): the method to use: "power", "torch". Default: "power".
        num_iterations (int): the number of iterations for "power" method. Default: 10.

    Returns:
        A Tensor of the leading eigenvector.
    """
    assert method in ["power", "torch"], f"Unsupported method: {method}."

    if method == "power":
        # power iteration algorithm
        leading_eig = torch.ones_like(mat[..., 0:1])  # (*, M, 1)
        leading_eig_last = leading_eig
        for i in range(num_iterations):
            leading_eig = torch.matmul(mat, leading_eig)
            leading_eig = F.normalize(leading_eig, p=2, dim=-2)
            if torch.allclose(leading_eig, leading_eig_last):
                break
            leading_eig_last = leading_eig
        leading_eig = leading_eig.squeeze(-1)
    else:
        eigenvalues, eigenvectors = torch.linalg.eigh(mat)
        leading_eig = eigenvectors[:, :, -1]

    return leading_eig
//  [python]
def ISTA_PN(I,basis,lambd,num_iter,eta=None, useMAGMA=True):
    # This is a positive-negative PyTorch-Ver ISTA solver
    # MAGMA uses CPU-GPU hybrid method to solve SVD problems, which is great for single task. When running multiple jobs, this flag should be turned off to leave the svd computation on only GPU.
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.symeig(torch.mm(basis,basis.t()),eigenvectors=False)[0])
            eta = 1./L
        else:
            eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()

    #Res = torch.zeros(I.size()).type(dtype)
    #ahat = torch.zeros(M,batch_size).type(dtype)
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in range(num_iter):
        ahat = ahat.add(eta * basis.t().mm(Res))
        ahat_sign = torch.sign(ahat)
        ahat.abs_()
        ahat.sub_(eta * lambd).clamp_(min = 0.)
        ahat.mul_(ahat_sign)
        Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def ISTA_PN(I,basis,lambd,num_iter,eta=None, useMAGMA=True):
    # This is a positive-negative PyTorch-Ver ISTA solver
    # MAGMA uses CPU-GPU hybrid method to solve SVD problems, which is great for single task. When running multiple jobs, this flag should be turned off to leave the svd computation on only GPU.
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.linalg.eigh(torch.mm(basis,basis.t()))[0])
            eta = 1./L
        else:
            eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()

    #Res = torch.zeros(I.size()).type(dtype)
    #ahat = torch.zeros(M,batch_size).type(dtype)
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in range(num_iter):
        ahat = ahat.add(eta * basis.t().mm(Res))
        ahat_sign = torch.sign(ahat)
        ahat.abs_()
        ahat.sub_(eta * lambd).clamp_(min = 0.)
        ahat.mul_(ahat_sign)
        Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def ISTA(I,basis,lambd,num_iter,eta=None, useMAGMA=True):
    # This is a positive-only PyTorch-Ver ISTA solver
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.symeig(torch.mm(basis,basis.t()),eigenvectors=False)[0])
            eta = 1./L
        else:
            eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()

    #Res = torch.zeros(I.size()).type(dtype)
    #ahat = torch.zeros(M,batch_size).type(dtype)
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in range(num_iter):
        ahat = ahat.add(eta * basis.t().mm(Res))
        ahat = ahat.sub(eta * lambd).clamp(min = 0.)
        Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def ISTA(I,basis,lambd,num_iter,eta=None, useMAGMA=True):
    # This is a positive-only PyTorch-Ver ISTA solver
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.linalg.eigh(torch.mm(basis,basis.t()),eigenvectors=False)[0])
            eta = 1./L
        else:
            eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()

    #Res = torch.zeros(I.size()).type(dtype)
    #ahat = torch.zeros(M,batch_size).type(dtype)
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in range(num_iter):
        ahat = ahat.add(eta * basis.t().mm(Res))
        ahat = ahat.sub(eta * lambd).clamp(min = 0.)
        Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def FISTA_adapt(I,basis,lambd,weights,num_iter,eta=None, useMAGMA=True):
    # This is a positive-only PyTorch-Ver ISTA solver
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.symeig(torch.mm(basis,basis.t()),eigenvectors=False)[0])
            eta = 1./L
        else:
            eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()

    tk_n = 1.
    tk = 1.
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)
    ahat_y = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in range(num_iter):
        tk = tk_n
        tk_n = (1+np.sqrt(1+4*tk**2))/2
        ahat_pre = ahat
        Res = I - torch.mm(basis,ahat_y)
        ahat_y = ahat_y.add(eta * basis.t().mm(Res))
        ahat = ahat_y.sub(eta * lambd * weights).clamp(min = 0.)
        ahat_y = ahat.add(ahat.sub(ahat_pre).mul((tk-1)/(tk_n)))
    Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def FISTA_adapt(I,basis,lambd,weights,num_iter,eta=None, useMAGMA=True):
    # This is a positive-only PyTorch-Ver ISTA solver
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.linalg.eigh(torch.mm(basis,basis.t()))[0])
            eta = 1./L
        else:
            eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()

    tk_n = 1.
    tk = 1.
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)
    ahat_y = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in range(num_iter):
        tk = tk_n
        tk_n = (1+np.sqrt(1+4*tk**2))/2
        ahat_pre = ahat
        Res = I - torch.mm(basis,ahat_y)
        ahat_y = ahat_y.add(eta * basis.t().mm(Res))
        ahat = ahat_y.sub(eta * lambd * weights).clamp(min = 0.)
        ahat_y = ahat.add(ahat.sub(ahat_pre).mul((tk-1)/(tk_n)))
    Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def ISTA_adapt(I,basis,lambd,weights,num_iter,eta=None,useMAGMA=True):
    # This is a PyTorch-Ver ISTA solver
    # The adaptive version ISTA will have a different shrinkage, determined by weights, on each dimensional.
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.symeig(torch.mm(basis,basis.t()),eigenvectors=False)[0])
            eta = 1./L
        else:
            eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()

    #Res = torch.zeros(I.size()).type(dtype)
    #ahat = torch.zeros(M,batch_size).type(dtype)
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in xrange(num_iter):
        ahat = ahat.add(eta * basis.t().mm(Res))
        ahat = ahat.sub_(eta * lambd * weights).clamp(min = 0.)
        Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def ISTA_adapt(I,basis,lambd,weights,num_iter,eta=None,useMAGMA=True):
    # This is a PyTorch-Ver ISTA solver
    # The adaptive version ISTA will have a different shrinkage, determined by weights, on each dimensional.
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.linalg.eigh(torch.mm(basis,basis.t()))[0])
            eta = 1./L
        else:
            eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()

    #Res = torch.zeros(I.size()).type(dtype)
    #ahat = torch.zeros(M,batch_size).type(dtype)
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in xrange(num_iter):
        ahat = ahat.add(eta * basis.t().mm(Res))
        ahat = ahat.sub_(eta * lambd * weights).clamp(min = 0.)
        Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def __compute_matrices(self, I, Io, alpha, beta):
        OD, ODhat = self.__convert_rgb2od(I, Io=Io, beta=beta)

        # compute eigenvectors
        _, eigvecs = torch.symeig(cov(ODhat.T), eigenvectors=True)
        eigvecs = eigvecs[:, [1, 2]]

        HE = self.__find_HE(ODhat, eigvecs, alpha)

        C = self.__find_concentration(OD, HE)
        maxC = torch.stack([percentile(C[0, :], 99), percentile(C[1, :], 99)])

        return HE, C, maxC
//  [python]
def __compute_matrices(self, I, Io, alpha, beta):
        OD, ODhat = self.__convert_rgb2od(I, Io=Io, beta=beta)

        # compute eigenvectors
        _, eigvecs = torch.linalg.eigh(cov(ODhat.T))
        eigvecs = eigvecs[:, [1, 2]]

        HE = self.__find_HE(ODhat, eigvecs, alpha)

        C = self.__find_concentration(OD, HE)
        maxC = torch.stack([percentile(C[0, :], 99), percentile(C[1, :], 99)])

        return HE, C, maxC
//  [python]
def ISTA_tempo(I,P,basis,D,tilde,lambd,num_iter,step_mod = 1.0,eta=None, useMAGMA=True):
    # This is a positive-only PyTorch-Ver ISTA solver
    # temporal regularization: we use pooling to regularize the sparse inference
    # objective: min ||I - Phi A||_2^2 + \tilde ||PAD||_2^2 + \lambda ||A||_1
    # \tilde is the temporal regularization coeff
    # \lambd is the sparse regularization coeff

    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.symeig(torch.mm(basis,basis.t()),eigenvectors=False)[0])
            eta = step_mod/L
        else:
            eta = step_mod/cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()
    DD = torch.mm(D,D.t())

    #Res = torch.zeros(I.size()).type(dtype)
    #ahat = torch.zeros(M,batch_size).type(dtype)
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in range(num_iter):
        ahat = ahat.add(eta * basis.t().mm(Res))
        ADD = torch.mm(ahat,DD)
        PADD = torch.mm(P,ADD)
        PPADD = torch.mm(P.t(),PADD)
        ahat = ahat.sub(tilde * eta * PPADD)
        ahat = ahat.sub_(eta * lambd).clamp(min = 0.)
        Res = I - torch.mm(basis,ahat)

    #For testing purpose:
    PAD = torch.mm(P,torch.mm(ahat,D))
    return ahat, Res, PAD
//  [python]
def ISTA_tempo(I,P,basis,D,tilde,lambd,num_iter,step_mod = 1.0,eta=None, useMAGMA=True):
    # This is a positive-only PyTorch-Ver ISTA solver
    # temporal regularization: we use pooling to regularize the sparse inference
    # objective: min ||I - Phi A||_2^2 + \tilde ||PAD||_2^2 + \lambda ||A||_1
    # \tilde is the temporal regularization coeff
    # \lambd is the sparse regularization coeff

    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
            L = torch.max(torch.linalg.eigh(torch.mm(basis,basis.t()),eigenvectors=False)[0])
            eta = step_mod/L
        else:
            eta = step_mod/cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
            eta = torch.from_numpy(eta.astype('float32')).cuda()
    DD = torch.mm(D,D.t())

    #Res = torch.zeros(I.size()).type(dtype)
    #ahat = torch.zeros(M,batch_size).type(dtype)
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)

    for t in range(num_iter):
        ahat = ahat.add(eta * basis.t().mm(Res))
        ADD = torch.mm(ahat,DD)
        PADD = torch.mm(P,ADD)
        PPADD = torch.mm(P.t(),PADD)
        ahat = ahat.sub(tilde * eta * PPADD)
        ahat = ahat.sub_(eta * lambd).clamp(min = 0.)
        Res = I - torch.mm(basis,ahat)

    #For testing purpose:
    PAD = torch.mm(P,torch.mm(ahat,D))
    return ahat, Res, PAD
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay
            if self.steps % self.Tf == 0:
                self.d_g[m], self.Q_g[m] = torch.symeig(
                    self.m_gg[m], eigenvectors=True)
                self.d_a[m], self.Q_a[m] = torch.symeig(
                    self.m_aa[m], eigenvectors=True)

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d' or classname == 'ConvTranspose2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data
            if self.Q_a[m].device != self.Q_g[m].device:
                self.Q_a[m]=self.Q_a[m].to(self.Q_g[m].device)
                self.d_a[m]=self.d_a[m].to(self.Q_g[m].device)

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()
            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            if p.requires_grad == False:
                continue
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()


        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            if p.requires_grad == False:
                continue
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
//  [python]
def step(self):
        # Add weight decay
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

        updates = {}
        for i, m in enumerate(self.modules):
            assert len(list(m.parameters())
                       ) == 1, "Can handle only one parameter at the moment"
            classname = m.__class__.__name__
            p = next(m.parameters())

            la = self.damping + self.weight_decay
            if self.steps % self.Tf == 0:
                self.d_g[m], self.Q_g[m] = torch.linalg.eigh(
                    self.m_gg[m])
                self.d_a[m], self.Q_a[m] = torch.linalg.eigh(
                    self.m_aa[m])

                self.d_a[m].mul_((self.d_a[m] > 1e-6).float())
                self.d_g[m].mul_((self.d_g[m] > 1e-6).float())

            if classname == 'Conv2d' or classname == 'ConvTranspose2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data
            if self.Q_a[m].device != self.Q_g[m].device:
                self.Q_a[m]=self.Q_a[m].to(self.Q_g[m].device)
                self.d_a[m]=self.d_a[m].to(self.Q_g[m].device)

            v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
            v2 = v1 / (
                self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + la)
            v = self.Q_g[m] @ v2 @ self.Q_a[m].t()
            v = v.view(p.grad.data.size())
            updates[p] = v

        vg_sum = 0
        for p in self.model.parameters():
            if p.requires_grad == False:
                continue
            v = updates[p]
            vg_sum += (v * p.grad.data * self.lr * self.lr).sum()


        nu = min(1, math.sqrt(self.kl_clip / vg_sum))

        for p in self.model.parameters():
            if p.requires_grad == False:
                continue
            v = updates[p]
            p.grad.data.copy_(v)
            p.grad.data.mul_(nu)

        self.optim.step()
        self.steps += 1
//  [python]
def test_symeig(self):
        def func(root, upper):
            x = 0.5 * (root + root.transpose(-2, -1))
            return torch.symeig(x, eigenvectors=True, upper=upper)

        def run_test(upper, dims):
            root = torch.rand(*dims, requires_grad=True)

            gradcheck(func, [root, upper])
            gradgradcheck(func, [root, upper])

            root = random_symmetric_matrix(dims[-1], *dims[:-2]).requires_grad_()
            w, v = root.symeig(eigenvectors=True)
            (w.sum() + v.sum()).backward()
            self.assertEqual(root.grad, root.grad.transpose(-1, -2))  # Check the gradient is symmetric

        for upper, dims in product([True, False], [(3, 3), (5, 3, 3), (4, 3, 2, 2)]):
            run_test(upper, dims)
//  [python]
def test_symeig(self):
        def func(root, upper):
            x = 0.5 * (root + root.transpose(-2, -1))
            return torch.linalg.eigh(x, UPLO='U' if upper else 'L')

        def run_test(upper, dims):
            root = torch.rand(*dims, requires_grad=True)

            gradcheck(func, [root, upper])
            gradgradcheck(func, [root, upper])

            root = random_symmetric_matrix(dims[-1], *dims[:-2]).requires_grad_()
            w, v = torch.linalg.eigh(root)
            (w.sum() + v.sum()).backward()
            self.assertEqual(root.grad, root.grad.transpose(-1, -2))  # Check the gradient is symmetric

        for upper, dims in product([True, False], [(3, 3), (5, 3, 3), (4, 3, 2, 2)]):
            run_test(upper, dims)
//  [python]
def test_batch_symeig_forward():
    xs = torch.randn(8, 4, 4).float().cuda()
    ws, vs = S.symeig(xs)
#    test=torch.symeig(xs[0])
    eigenValues, eigenVectors = torch.symeig(xs[0],eigenvectors=True)
    for i in range(xs.shape[0]):
        w, v = S.symeig(xs[i])
        torch.testing.assert_allclose(ws[i], w)
        torch.testing.assert_allclose(vs[i], v)
//  [python]
def test_batch_symeig_forward():
    xs = torch.randn(8, 4, 4).float().cuda()
    ws, vs = S.symeig(xs)
#    test=torch.symeig(xs[0])
    eigenValues, eigenVectors = torch.linalg.eigh(xs[0])
    for i in range(xs.shape[0]):
        w, v = S.symeig(xs[i])
        torch.testing.assert_allclose(ws[i], w)
        torch.testing.assert_allclose(vs[i], v)
//  [python]
def get_Laplace_from_pc(ori_pc):
    with torch.no_grad():
        point_mat = ori_pc.unsqueeze(1) - ori_pc.unsqueeze(0)
        A = torch.exp(-torch.sum(point_mat.square(), dim=2))

        mask = torch.zeros_like(A)
        idx = knn(ori_pc, 30)
        for i in range(idx.shape[0]):
            for jid in range(idx.shape[1]):
                j = idx[i, jid]
                mask[i, j] = 1
                mask[j, i] = 1

        A = A * mask
        D = torch.diag(torch.sum(A, dim=1))
        L = D - A
        # D_sqtinv = torch.diag(1.0 / torch.sqrt(torch.sum(A, dim=1)))
        # L = torch.eye(ori_pc.shape[0]).to(A.device) - torch.matmul(torch.matmul(D_sqtinv, A), D_sqtinv)
        e, v = torch.symeig(L, eigenvectors=True)
    return e, v
//  [python]
def get_Laplace_from_pc(ori_pc):
    with torch.no_grad():
        point_mat = ori_pc.unsqueeze(1) - ori_pc.unsqueeze(0)
        A = torch.exp(-torch.sum(point_mat.square(), dim=2))

        mask = torch.zeros_like(A)
        idx = knn(ori_pc, 30)
        for i in range(idx.shape[0]):
            for jid in range(idx.shape[1]):
                j = idx[i, jid]
                mask[i, j] = 1
                mask[j, i] = 1

        A = A * mask
        D = torch.diag(torch.sum(A, dim=1))
        L = D - A
        # D_sqtinv = torch.diag(1.0 / torch.sqrt(torch.sum(A, dim=1)))
        # L = torch.eye(ori_pc.shape[0]).to(A.device) - torch.matmul(torch.matmul(D_sqtinv, A), D_sqtinv)
        e, v = torch.linalg.eigh(L)
    return e, v
//  [python]
def reyi_entropy(x,sigma):
    alpha = 1.01
    k = calculate_gram_mat(x,sigma)
    k = k/torch.trace(k) 
    eigv = torch.abs(torch.symeig(k, eigenvectors=True)[0])
    eig_pow = eigv**alpha
    entropy = (1/(1-alpha))*torch.log2(torch.sum(eig_pow))
    return entropy
//  [python]
def reyi_entropy(x,sigma):
    alpha = 1.01
    k = calculate_gram_mat(x,sigma)
    k = k/torch.trace(k) 
    eigv = torch.abs(torch.linalg.eigh(k)[0])
    eig_pow = eigv**alpha
    entropy = (1/(1-alpha))*torch.log2(torch.sum(eig_pow))
    return entropy
//  [python]
def FISTA(I,basis,lambd,num_iter,sigma,eta=None, useMAGMA=True, plot=False):
    # This is a positive-only PyTorch-Ver ISTA solver
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
#             L = torch.max(torch.symeig(torch.mm(basis,basis.t()),eigenvectors=False)[0])/sigma**2
            jittered_basis = basis + torch.rand(basis.size()).cuda()*1.0e-15
            L = torch.from_numpy(np.array([np.nanmax(torch.symeig(torch.mm(jittered_basis,jittered_basis.t()),eigenvectors=False)[0].cpu().numpy())])).cuda()/sigma**2
            eta = 1./L
#         else: # Need to change with sigma
#             eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
#             eta = torch.from_numpy(eta.astype('float32')).cuda()

    tk_n = 1.
    tk = 1.
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)
    ahat_y = torch.cuda.FloatTensor(M,batch_size).fill_(0)
    
    if plot:
        fig = plt.figure(figsize=(10,5))
    for t in range(num_iter):
        tk = tk_n
        tk_n = (1+np.sqrt(1+4*tk**2))/2
        ahat_pre = ahat
        Res = I - torch.mm(basis,ahat_y)
#         print("ahat_y", torch.isnan(ahat_y).any())
#         print("RES", torch.isnan(Res).any())
#         print("eta", torch.isnan(eta).any())
        ahat_y = ahat_y.add(eta * basis.t().mm(Res)/sigma**2)
#         print(torch.isnan(ahat_y).any())
        ahat = ahat_y.sub(eta * lambd).clamp(min = 0.)
#         print(torch.isnan(ahat).any())
        ahat_y = ahat.add(ahat.sub(ahat_pre).mul((tk-1)/(tk_n)))
#         print(torch.isnan(ahat_y).any())
        
        # Plot snr and alpha
        if plot and t % 1 == 0:
            snr = ((I**2).sum(dim=0)/(Res**2).sum(dim=0)).mean()
            fig.add_subplot(1,2,1)
            plt.scatter(t, snr.cpu().numpy())
            fig.add_subplot(1,2,2)
            plt.scatter(np.ones(ahat.size()[0])*t, ahat[:,0].cpu().numpy(), c=np.arange(ahat.size()[0]))
            display.clear_output(wait=True)
            display.display(plt.gcf())
    if plot:
        display.clear_output()
    Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def FISTA(I,basis,lambd,num_iter,sigma,eta=None, useMAGMA=True, plot=False):
    # This is a positive-only PyTorch-Ver ISTA solver
    dtype = basis.type()
    batch_size=I.size(1)
    M = basis.size(1)
    if eta is None:
        if useMAGMA:
#             L = torch.max(torch.symeig(torch.mm(basis,basis.t()),eigenvectors=False)[0])/sigma**2
            jittered_basis = basis + torch.rand(basis.size()).cuda()*1.0e-15
            L = torch.from_numpy(np.array([np.nanmax(torch.linalg.eigh(torch.mm(jittered_basis,jittered_basis.t()))[0].cpu().numpy())])).cuda()/sigma**2
            eta = 1./L
#         else: # Need to change with sigma
#             eta = 1./cp.linalg.eigvalsh(cp.asarray(torch.mm(basis,basis.t()).cpu().numpy())).max().get().reshape(1)
#             eta = torch.from_numpy(eta.astype('float32')).cuda()

    tk_n = 1.
    tk = 1.
    Res = torch.cuda.FloatTensor(I.size()).fill_(0)
    ahat = torch.cuda.FloatTensor(M,batch_size).fill_(0)
    ahat_y = torch.cuda.FloatTensor(M,batch_size).fill_(0)
    
    if plot:
        fig = plt.figure(figsize=(10,5))
    for t in range(num_iter):
        tk = tk_n
        tk_n = (1+np.sqrt(1+4*tk**2))/2
        ahat_pre = ahat
        Res = I - torch.mm(basis,ahat_y)
#         print("ahat_y", torch.isnan(ahat_y).any())
#         print("RES", torch.isnan(Res).any())
#         print("eta", torch.isnan(eta).any())
        ahat_y = ahat_y.add(eta * basis.t().mm(Res)/sigma**2)
#         print(torch.isnan(ahat_y).any())
        ahat = ahat_y.sub(eta * lambd).clamp(min = 0.)
#         print(torch.isnan(ahat).any())
        ahat_y = ahat.add(ahat.sub(ahat_pre).mul((tk-1)/(tk_n)))
#         print(torch.isnan(ahat_y).any())
        
        # Plot snr and alpha
        if plot and t % 1 == 0:
            snr = ((I**2).sum(dim=0)/(Res**2).sum(dim=0)).mean()
            fig.add_subplot(1,2,1)
            plt.scatter(t, snr.cpu().numpy())
            fig.add_subplot(1,2,2)
            plt.scatter(np.ones(ahat.size()[0])*t, ahat[:,0].cpu().numpy(), c=np.arange(ahat.size()[0]))
            display.clear_output(wait=True)
            display.display(plt.gcf())
    if plot:
        display.clear_output()
    Res = I - torch.mm(basis,ahat)
    return ahat, Res
//  [python]
def test_symeig():
    # NOTE need pytorch 0.5.0 or 1.0
    a = torch.tensor([[ 1.96,  0.00,  0.00,  0.00,  0.00],
                      [-6.49,  3.80,  0.00,  0.00,  0.00],
                      [-0.47, -6.39,  4.17,  0.00,  0.00],
                      [-7.20,  1.50, -1.51,  5.70,  0.00],
                      [-0.65, -6.34,  2.67,  1.80, -7.10]]).t()
    a.requires_grad = True
    w, v = torch.symeig(a, eigenvectors=True)
    v.sum().backward()
    print(v.grad)
//  [python]
def test_symeig():
    # NOTE need pytorch 0.5.0 or 1.0
    a = torch.tensor([[ 1.96,  0.00,  0.00,  0.00,  0.00],
                      [-6.49,  3.80,  0.00,  0.00,  0.00],
                      [-0.47, -6.39,  4.17,  0.00,  0.00],
                      [-7.20,  1.50, -1.51,  5.70,  0.00],
                      [-0.65, -6.34,  2.67,  1.80, -7.10]]).t()
    a.requires_grad = True
    w, v = torch.linalg.eigh(a)
    v.sum().backward()
    print(v.grad)
//  [python]
def Construct_LRF(local_cor, local_weight=None, ref_point=None):
    device = local_cor.device
    B, N, _ = local_cor.shape
    cov_matrix = torch.matmul(local_cor.transpose(-1, -2), local_cor) if local_weight is None \
        else Variable(torch.matmul(local_cor.transpose(-1, -2), local_cor * local_weight), requires_grad=True)
    Z_axis = torch.symeig(cov_matrix, eigenvectors=True)[1][:, :, 0]  # Covariance matrix decomposition
    mask = (torch.sum(-Z_axis * ref_point, dim=1) < 0).float().unsqueeze(1)
    Z_axis = Z_axis * (1 - mask) - Z_axis * mask

    return Z_axis
//  [python]
def Construct_LRF(local_cor, local_weight=None, ref_point=None):
    device = local_cor.device
    B, N, _ = local_cor.shape
    cov_matrix = torch.matmul(local_cor.transpose(-1, -2), local_cor) if local_weight is None \
        else Variable(torch.matmul(local_cor.transpose(-1, -2), local_cor * local_weight), requires_grad=True)
    Z_axis = torch.linalg.eigh(cov_matrix)[1][:, :, 0]  # Covariance matrix decomposition
    mask = (torch.sum(-Z_axis * ref_point, dim=1) < 0).float().unsqueeze(1)
    Z_axis = Z_axis * (1 - mask) - Z_axis * mask

    return Z_axis
//  [python]
def Rotation_invariance_weight(local_cor):
    device = local_cor.device
    B, N, _ = local_cor.shape
    distance=(local_cor[:,:,0]**2+local_cor[:,:,1]**2+local_cor[:,:,2]**2)**0.5
    distance_max,_=torch.max(distance,dim=1)
    weight=distance_max.view(B,1,1).expand(B,N,1)-distance.view(B,N,1)
    weight_2=weight/torch.sum(weight,dim=1,keepdim=True).expand(B,N,1)
    cov_matrix=torch.matmul(local_cor.transpose(-1, -2), local_cor * weight_2.expand(B,N,3))
    Z_axis = torch.symeig(cov_matrix, eigenvectors=True)[1][:, :, 0]  #计算特征值（和向量），这里得到他的最大的特征向量

    return Z_axis
//  [python]
def Rotation_invariance_weight(local_cor):
    device = local_cor.device
    B, N, _ = local_cor.shape
    distance=(local_cor[:,:,0]**2+local_cor[:,:,1]**2+local_cor[:,:,2]**2)**0.5
    distance_max,_=torch.max(distance,dim=1)
    weight=distance_max.view(B,1,1).expand(B,N,1)-distance.view(B,N,1)
    weight_2=weight/torch.sum(weight,dim=1,keepdim=True).expand(B,N,1)
    cov_matrix=torch.matmul(local_cor.transpose(-1, -2), local_cor * weight_2.expand(B,N,3))
    Z_axis = torch.linalg.eigh(cov_matrix, eigenvectors=True)[1][:, :, 0]  #计算特征值（和向量），这里得到他的最大的特征向量

    return Z_axis
//  [python]
def joint_entropy(x,y,s_x,s_y):
    alpha = 1.01
    x = calculate_gram_mat(x,s_x)
    y = calculate_gram_mat(y,s_y)
    k = torch.mul(x,y)
    k = k/torch.trace(k)
    eigv = torch.abs(torch.symeig(k, eigenvectors=True)[0])
    eig_pow =  eigv**alpha
    entropy = (1/(1-alpha))*torch.log2(torch.sum(eig_pow))

    return entropy
//  [python]
def joint_entropy(x,y,s_x,s_y):
    alpha = 1.01
    x = calculate_gram_mat(x,s_x)
    y = calculate_gram_mat(y,s_y)
    k = torch.mul(x,y)
    k = k/torch.trace(k)
    eigv = torch.abs(torch.linalg.eigh(k)[0])
    eig_pow =  eigv**alpha
    entropy = (1/(1-alpha))*torch.log2(torch.sum(eig_pow))

    return entropy
//  [python]
def test_symeig_positive_lanczos(self, D=8):
        "Compare `eff_dim.symeig_lanczos` with `torch.symeig`"
        M = torch.randn(D, D)
        M = M@M.t()
        ref_vals, ref_vecs = torch.symeig(M, eigenvectors=True)

        # All values/vectors:
        vals, vecs = eff_dim.symeig_positive_lanczos(M, n_eigs=-1, vecs=True)
        assert torch.allclose(ref_vals, vals, rtol=1e-4)
        _assert_eigvecs_allclose(ref_vecs, vecs, rtol=1e-3)
//  [python]
def test_symeig_positive_lanczos(self, D=8):
        "Compare `eff_dim.symeig_lanczos` with `torch.linalg.eigh`"
        M = torch.randn(D, D)
        M = M@M.t()
        ref_vals, ref_vecs = torch.linalg.eigh(M)

        # All values/vectors:
        vals, vecs = eff_dim.symeig_positive_lanczos(M, n_eigs=-1, vecs=True)
        assert torch.allclose(ref_vals, vals, rtol=1e-4)
        _assert_eigvecs_allclose(ref_vecs, vecs, rtol=1e-3)