jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, random_seed_pair=(None, None)):
        self.n_vertices = np.array([2] * LABS_DIM)
        self.suggested_init = torch.empty(0).long()
        self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed_pair[1])], dim=0)
        self.adjacency_mat = []
        self.fourier_freq = []
        self.fourier_basis = []
        self.random_seed_info = 'R'.join([str(random_seed_pair[i]).zfill(4) if random_seed_pair[i] is not None else 'None' for i in range(2)])
        for i in range(len(self.n_vertices)):
            n_v = self.n_vertices[i]
            adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
            self.adjacency_mat.append(adjmat)
            laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
            eigval, eigvec = torch.symeig(laplacian, eigenvectors=True)
            self.fourier_freq.append(eigval)
            self.fourier_basis.append(eigvec)
//  [python]
def __init__(self, random_seed_pair=(None, None)):
        self.n_vertices = np.array([2] * LABS_DIM)
        self.suggested_init = torch.empty(0).long()
        self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed_pair[1])], dim=0)
        self.adjacency_mat = []
        self.fourier_freq = []
        self.fourier_basis = []
        self.random_seed_info = 'R'.join([str(random_seed_pair[i]).zfill(4) if random_seed_pair[i] is not None else 'None' for i in range(2)])
        for i in range(len(self.n_vertices)):
            n_v = self.n_vertices[i]
            adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
            self.adjacency_mat.append(adjmat)
            laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
            eigval, eigvec = torch.linalg.eigh(laplacian)
            self.fourier_freq.append(eigval)
            self.fourier_basis.append(eigvec)
//  [python]
def linalg_bench():
    if np.__config__.get_info("lapack_mkl_info"):
        print("MKL version", get_mkl_version())
    else:
        print("not using MKL")

    print("PyTorch version", torch.version.__version__)

    print("Scipy version: ", scipy.version.full_version)
    print("Numpy version: ", np.version.full_version)

    times = {}
    dimlist = [768, 2500, 2000, 1500, 1000, 500]
    for device in ['cpu', 'cuda']:
        for d in dimlist*2:
            print(f"{d}-by-{d} matrix: ", device)
            n = 10000
            assert n > 2*d   # to prevent singularity
            X = u.from_numpy(np.random.random((d, 10000)))
            Y = u.from_numpy(np.random.random((d, 10000)))
            H = X @ X.t()
            S = Y @ Y.t()
            if torch.cuda.is_available():
                H = H.to(device)
                S = S.to(device)

            with timeit(f"symeig"):
                result = torch.symeig(H, eigenvectors=False)
            times.setdefault(d, []).append(last_time)
            
            with timeit(f"symeig"):
                result = torch.symeig(H, eigenvectors=True)
            times.setdefault(d, []).append(last_time)

            with timeit(f"svd"):
                result = torch.svd(H)
            times.setdefault(d, []).append(last_time)
    for d in times:
        print(f"{d}-by-{d}: {u.format_list(times[d][len(times[d])//2:])}")
//  [python]
def linalg_bench():
    if np.__config__.get_info("lapack_mkl_info"):
        print("MKL version", get_mkl_version())
    else:
        print("not using MKL")

    print("PyTorch version", torch.version.__version__)

    print("Scipy version: ", scipy.version.full_version)
    print("Numpy version: ", np.version.full_version)

    times = {}
    dimlist = [768, 2500, 2000, 1500, 1000, 500]
    for device in ['cpu', 'cuda']:
        for d in dimlist*2:
            print(f"{d}-by-{d} matrix: ", device)
            n = 10000
            assert n > 2*d   # to prevent singularity
            X = u.from_numpy(np.random.random((d, 10000)))
            Y = u.from_numpy(np.random.random((d, 10000)))
            H = X @ X.t()
            S = Y @ Y.t()
            if torch.cuda.is_available():
                H = H.to(device)
                S = S.to(device)

            with timeit(f"symeig"):
                result = torch.linalg.eigh(H, eigenvectors=False)
            times.setdefault(d, []).append(last_time)
            
            with timeit(f"symeig"):
                result = torch.linalg.eigh(H, eigenvectors=True)
            times.setdefault(d, []).append(last_time)

            with timeit(f"svd"):
                result = torch.linalg.svd(H)
            times.setdefault(d, []).append(last_time)
    for d in times:
        print(f"{d}-by-{d}: {u.format_list(times[d][len(times[d])//2:])}")
//  [python]
def alpha_update_FISTA(y, psi, w, omega, alpha, sigma, lamb, lamb2, steps, k=None, m=None, eta_alpha=0.001, N=100, modified=False, plot=False, map_est=False, adaptive=True, adaptive_steps=False):
    D, K = psi.size()[0], psi.size()[1]
    tk_n, tk = 1.0, 1.0
    prev_alpha, alpha_y = alpha.clone(), alpha.clone()
    if adaptive:
        wpsi = torch.mm(w.t(),psi)
        lipschitz = 1.5*torch.max(torch.symeig(torch.mm(wpsi.t(),wpsi),eigenvectors=False)[0])/sigma**2
        eta_alpha = 1.0/lipschitz
#         print(1/eta_alpha)
        if adaptive_steps:
            steps = int(21*torch.sqrt(0.001/eta_alpha))-1
#             print(eta_alpha, int(21*torch.sqrt(0.001/eta_alpha))-1)
    if plot:
        fig = plt.figure(figsize=(10,5))
    for t in range(steps):
        # From here till end excluding compute_aux_var (0.0033s)
        tk = tk_n
        tk_n = (1+np.sqrt(1+4*tk**2))/2
        alpha_pre = alpha.clone()
        x, u, v, k_hat, m_hat, q, res, tres = compute_aux_var(y, psi, w, alpha_y, omega, sigma, k=k, m=m, N=N, map_est=map_est) # (0.0032s)
        if modified:
            dalpha = eta(1, eta_0=eta_alpha)*grad_alpha_modified(tres, psi, alpha_y, prev_alpha, sigma, lamb, lamb2)
        else:
            dalpha = eta(1, eta_0=eta_alpha)*grad_alpha(tres, psi, alpha_y, prev_alpha, sigma, lamb, lamb2)
        alpha = (alpha_y + dalpha).clamp(min=0.)
        alpha_y = alpha + (tk-1)/tk_n * (alpha - alpha_pre)
        
        # Plot loss and alpha
        if plot and t % 1 == 0:
            nll = -log_likelihood(x, y, u, v, psi, alpha, k_hat, m_hat, omega, lamb, lamb2, sigma, mean=True, N=N, modified=modified)
            fig.add_subplot(1,2,1)
            plt.scatter(t, nll)
            fig.add_subplot(1,2,2)
            plt.scatter(np.ones(alpha.size()[1])*t, alpha[0].cpu().numpy(), c=np.arange(K))
            display.clear_output(wait=True)
            display.display(plt.gcf())
    if plot:
        display.clear_output()
    return alpha
//  [python]
def alpha_update_FISTA(y, psi, w, omega, alpha, sigma, lamb, lamb2, steps, k=None, m=None, eta_alpha=0.001, N=100, modified=False, plot=False, map_est=False, adaptive=True, adaptive_steps=False):
    D, K = psi.size()[0], psi.size()[1]
    tk_n, tk = 1.0, 1.0
    prev_alpha, alpha_y = alpha.clone(), alpha.clone()
    if adaptive:
        wpsi = torch.mm(w.t(),psi)
        lipschitz = 1.5*torch.max(torch.linalg.eigh(torch.mm(wpsi.t(),wpsi))[0])/sigma**2
        eta_alpha = 1.0/lipschitz
#         print(1/eta_alpha)
        if adaptive_steps:
            steps = int(21*torch.sqrt(0.001/eta_alpha))-1
#             print(eta_alpha, int(21*torch.sqrt(0.001/eta_alpha))-1)
    if plot:
        fig = plt.figure(figsize=(10,5))
    for t in range(steps):
        # From here till end excluding compute_aux_var (0.0033s)
        tk = tk_n
        tk_n = (1+np.sqrt(1+4*tk**2))/2
        alpha_pre = alpha.clone()
        x, u, v, k_hat, m_hat, q, res, tres = compute_aux_var(y, psi, w, alpha_y, omega, sigma, k=k, m=m, N=N, map_est=map_est) # (0.0032s)
        if modified:
            dalpha = eta(1, eta_0=eta_alpha)*grad_alpha_modified(tres, psi, alpha_y, prev_alpha, sigma, lamb, lamb2)
        else:
            dalpha = eta(1, eta_0=eta_alpha)*grad_alpha(tres, psi, alpha_y, prev_alpha, sigma, lamb, lamb2)
        alpha = (alpha_y + dalpha).clamp(min=0.)
        alpha_y = alpha + (tk-1)/tk_n * (alpha - alpha_pre)
        
        # Plot loss and alpha
        if plot and t % 1 == 0:
            nll = -log_likelihood(x, y, u, v, psi, alpha, k_hat, m_hat, omega, lamb, lamb2, sigma, mean=True, N=N, modified=modified)
            fig.add_subplot(1,2,1)
            plt.scatter(t, nll)
            fig.add_subplot(1,2,2)
            plt.scatter(np.ones(alpha.size()[1])*t, alpha[0].cpu().numpy(), c=np.arange(K))
            display.clear_output(wait=True)
            display.display(plt.gcf())
    if plot:
        display.clear_output()
    return alpha
//  [python]
def _symeig(self, eigenvectors: bool = False) -> Tuple[torch.Tensor, Optional["LinearOperator"]]:
        """Method that allows implementing special-cased symeig computation. Should not be called directly"""
        from .dense_linear_operator import DenseLinearOperator

        dtype = self.dtype  # perform decomposition in double precision for numerical stability
        # TODO: Use fp64 registry once #1213 is addressed
        evals, evecs = torch.symeig(self.to_dense().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 = DenseLinearOperator(evecs.to(dtype=dtype))
        else:
            evecs = None
        return evals, evecs
//  [python]
def _symeig(self, eigenvectors: bool = False) -> Tuple[torch.Tensor, Optional["LinearOperator"]]:
        """Method that allows implementing special-cased symeig computation. Should not be called directly"""
        from .dense_linear_operator import DenseLinearOperator

        dtype = self.dtype  # perform decomposition in double precision for numerical stability
        # TODO: Use fp64 registry once #1213 is addressed
        evals, evecs = torch.linalg.eigh(self.to_dense().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 = DenseLinearOperator(evecs.to(dtype=dtype))
        else:
            evecs = None
        return evals, evecs
//  [python]
def eigh(tensor):
        """Legacy only, deprecated from PyTorch 1.8.0"""
        return torch.symeig(tensor, eigenvectors=True)
//  [python]
def eigh(tensor):
        """Legacy only, deprecated from PyTorch 1.8.0"""
        return torch.linalg.eigh(tensor)
//  [python]
def __init__(self, random_seed_pair=(None, None)):
		citywise_satisfaction, tourism_attractions_similarity, transportation_type, travel_time, cost, delay_prob = generate_travel_plan_problem(TRAVEL_N_CITIES, random_seed_pair[0])
		self.n_vertices = [TRAVEL_N_CITIES] + list(number_of_edges(transportation_type))
		self.suggested_init = torch.empty(0).long()
		self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed_pair[1]).long()], dim=0)
		self.adjacency_mat = []
		self.fourier_freq = []
		self.fourier_basis = []
		self.random_seed_info = 'R'.join([str(random_seed_pair[h]).zfill(4) if random_seed_pair[h] is not None else 'None' for h in range(2)])
		for i in range(len(self.n_vertices)):
			n_v = self.n_vertices[i]
			adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
			self.adjacency_mat.append(adjmat)
			laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
			eigval, eigvec = torch.symeig(laplacian, eigenvectors=True)
			self.fourier_freq.append(eigval)
			self.fourier_basis.append(eigvec)
		self.citywise_satisfaction = citywise_satisfaction
		self.tourism_attractions_similarity = tourism_attractions_similarity
		self.transportation_type = transportation_type
		self.travel_time = travel_time
		self.cost = cost
		self.delay_prob = delay_prob
//  [python]
def __init__(self, random_seed_pair=(None, None)):
		citywise_satisfaction, tourism_attractions_similarity, transportation_type, travel_time, cost, delay_prob = generate_travel_plan_problem(TRAVEL_N_CITIES, random_seed_pair[0])
		self.n_vertices = [TRAVEL_N_CITIES] + list(number_of_edges(transportation_type))
		self.suggested_init = torch.empty(0).long()
		self.suggested_init = torch.cat([self.suggested_init, sample_init_points(self.n_vertices, 20 - self.suggested_init.size(0), random_seed_pair[1]).long()], dim=0)
		self.adjacency_mat = []
		self.fourier_freq = []
		self.fourier_basis = []
		self.random_seed_info = 'R'.join([str(random_seed_pair[h]).zfill(4) if random_seed_pair[h] is not None else 'None' for h in range(2)])
		for i in range(len(self.n_vertices)):
			n_v = self.n_vertices[i]
			adjmat = torch.diag(torch.ones(n_v - 1), -1) + torch.diag(torch.ones(n_v - 1), 1)
			self.adjacency_mat.append(adjmat)
			laplacian = torch.diag(torch.sum(adjmat, dim=0)) - adjmat
			eigval, eigvec = torch.linalg.eigh(laplacian)
			self.fourier_freq.append(eigval)
			self.fourier_basis.append(eigvec)
		self.citywise_satisfaction = citywise_satisfaction
		self.tourism_attractions_similarity = tourism_attractions_similarity
		self.transportation_type = transportation_type
		self.travel_time = travel_time
		self.cost = cost
		self.delay_prob = delay_prob
//  [python]
def batch_symeig(X):
  # it is much faster to run symeig on CPU
  X = X.cpu()
  b, d, _ = X.size()
  bv = X.new(b, d, d)
  for batch_idx in range(X.shape[0]):
    e, v = torch.symeig(X[batch_idx, :, :].squeeze(), True)
    bv[batch_idx, :, :] = v
  bv = bv.cuda()
  return bv
//  [python]
def batch_symeig(X):
  # it is much faster to run symeig on CPU
  X = X.cpu()
  b, d, _ = X.size()
  bv = X.new(b, d, d)
  for batch_idx in range(X.shape[0]):
    e, v = torch.linalg.eigh(X[batch_idx, :, :].squeeze())
    bv[batch_idx, :, :] = v
  bv = bv.cuda()
  return bv
//  [python]
def cal_leading_eigenvector(self, M, method='power'):
        if method == 'power':
            # power iteration algorithm
            leading_eig = torch.ones_like(M[:, :, 0:1])
            leading_eig_last = leading_eig
            for i in range(self.num_iterations):
                leading_eig = torch.bmm(M, leading_eig)
                leading_eig = leading_eig / (torch.norm(leading_eig, dim=1, keepdim=True) + 1e-6)
                if torch.allclose(leading_eig, leading_eig_last):
                    break
                leading_eig_last = leading_eig
            leading_eig = leading_eig.squeeze(-1)
            return leading_eig
        elif method == 'eig':  # cause NaN during back-prop
            e, v = torch.symeig(M, eigenvectors=True)
            leading_eig = v[:, :, -1]
            return leading_eig
        else:
            exit(-1)
//  [python]
def cal_leading_eigenvector(self, M, method='power'):
        if method == 'power':
            # power iteration algorithm
            leading_eig = torch.ones_like(M[:, :, 0:1])
            leading_eig_last = leading_eig
            for i in range(self.num_iterations):
                leading_eig = torch.bmm(M, leading_eig)
                leading_eig = leading_eig / (torch.norm(leading_eig, dim=1, keepdim=True) + 1e-6)
                if torch.allclose(leading_eig, leading_eig_last):
                    break
                leading_eig_last = leading_eig
            leading_eig = leading_eig.squeeze(-1)
            return leading_eig
        elif method == 'eig':  # cause NaN during back-prop
            e, v = torch.linalg.eigh(M)
            leading_eig = v[:, :, -1]
            return leading_eig
        else:
            exit(-1)
//  [python]
def stn(self, x, u):
        # A_vec = self.fc_stn(u)
        # print(x.reshape(-1, 28 * 28).size(),u.size())
        A_vec = self.fc_stn(torch.cat([u.unsqueeze(1), x.reshape(-1, 28 * 28)], 1))
        A = convert_Avec_to_A(A_vec)
        _, evs = torch.symeig(A, eigenvectors=True)
        tcos, tsin = evs[:, 0:1, 0:1], evs[:, 1:2, 0:1]

        self.theta_angle = torch.atan2(tsin[:, 0, 0], tcos[:, 0, 0])

        # clock-wise rotate theta
        theta_0 = torch.cat([tcos, tsin, tcos * 0], 2)
        theta_1 = torch.cat([-tsin, tcos, tcos * 0], 2)
        theta = torch.cat([theta_0, theta_1], 1)

        grid = F.affine_grid(theta, x.size())
        x = F.grid_sample(x, grid)
        return x
//  [python]
def stn(self, x, u):
        A_vec = self.fc_stn(torch.cat([u.unsqueeze(1), x.reshape(-1, 28 * 28)], 1))
        A = convert_Avec_to_A(A_vec)
        _, evs = torch.linalg.eigh(A)
        tcos, tsin = evs[:, 0:1, 0:1], evs[:, 1:2, 0:1]

        self.theta_angle = torch.atan2(tsin[:, 0, 0], tcos[:, 0, 0])

        theta_0 = torch.cat([tcos, tsin, tcos * 0], 2)
        theta_1 = torch.cat([-tsin, tcos, tcos * 0], 2)
        theta = torch.cat([theta_0, theta_1], 1)

        grid = F.affine_grid(theta, x.size())
        x = F.grid_sample(x, grid)
        return x
//  [python]
def forward(self, P0, P1):
		self.p1, self.V = torch.symeig(P1, eigenvectors = True)
		self.lnp1 = torch.log(self.p1)
		xent = - torch.sum(torch.sum(self.V * P0.mm(self.V), 0)*self.lnp1, 0)
		self.save_for_backward(P0, P1)
		return xent
//  [python]
def forward(self, P0, P1):
		self.p1, self.V = torch.linalg.eigh(P1)
		self.lnp1 = torch.log(self.p1)
		xent = - torch.sum(torch.sum(self.V * P0.mm(self.V), 0)*self.lnp1, 0)
		self.save_for_backward(P0, P1)
		return xent
//  [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 sqrtm(A):
    P, V = torch.symeig(A, eigenvectors=True)
    dtype_cond = {torch.float32: 1e-6, torch.float64: 1e-10}
    cond = dtype_cond[A.dtype]
    cutoff = (abs(P) > cond * torch.max(abs(P)))

    D = torch.sqrt(P[cutoff])
    S = torch.diag(D)
    V = V[:, cutoff]
    B = V @ S @ V.t()
    return B
//  [python]
def sqrtm(A):
    P, V = torch.linalg.eigh(A)
    dtype_cond = {torch.float32: 1e-6, torch.float64: 1e-10}
    cond = dtype_cond[A.dtype]
    cutoff = (abs(P) > cond * torch.max(abs(P)))

    D = torch.sqrt(P[cutoff])
    S = torch.diag(D)
    V = V[:, cutoff]
    B = V @ S @ V.t()
    return B
//  [python]
def ComputePCA(xNormalized):    
    n = xNormalized.shape[1]
    coefficient = 1/(n-1)
    covMatrix = torch.matmul(torch.transpose(xNormalized, 0, 1),xNormalized).to(args.deviceName)
    covMatrix = torch.mul(covMatrix,coefficient)
    torch.distributed.all_reduce(covMatrix,op=dist.ReduceOp.SUM,async_op=False) 
    eigenValues, eigenVectors = torch.symeig(covMatrix, eigenvectors=True)
    np.savetxt("eigenValues.csv",eigenValues.cpu().numpy(), delimiter=",") 
    np.savetxt("eigenVectors.csv",eigenVectors.cpu().numpy(), delimiter=",") 
    return eigenVectors
//  [python]
def ComputePCA(xNormalized):    
    n = xNormalized.shape[1]
    coefficient = 1/(n-1)
    covMatrix = torch.matmul(torch.transpose(xNormalized, 0, 1),xNormalized).to(args.deviceName)
    covMatrix = torch.mul(covMatrix,coefficient)
    torch.distributed.all_reduce(covMatrix,op=dist.ReduceOp.SUM,async_op=False) 
    eigenValues, eigenVectors = torch.linalg.eigh(covMatrix)
    np.savetxt("eigenValues.csv",eigenValues.cpu().numpy(), delimiter=",") 
    np.savetxt("eigenVectors.csv",eigenVectors.cpu().numpy(), delimiter=",") 
    return eigenVectors
//  [python]
def reyi_entropy(x, sigma):
    alpha = 1.01
    print("pairwise distance")
    k = calculate_gram_mat(x, sigma)
    k = k / torch.trace(k)
    print("compute eigv")
    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
    print("pairwise distance")
    k = calculate_gram_mat(x, sigma)
    k = k / torch.trace(k)
    print("compute eigv")
    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(self, device, dtype):
        from common_utils import random_symmetric_matrix

        def run_test(dims, eigenvectors, upper):
            x = random_symmetric_matrix(*dims, dtype=dtype, device=device)
            oute = torch.empty(dims[1:] + dims[:1], dtype=dtype, device=device)
            outv = torch.empty(dims[1:] + dims[:1] * 2, dtype=dtype, device=device)
            torch.symeig(x, eigenvectors=eigenvectors, upper=upper, out=(oute, outv))

            if eigenvectors:
                x_recon = torch.matmul(torch.matmul(outv, torch.diag_embed(oute)), outv.transpose(-2, -1))
                self.assertEqual(x, x_recon, 1e-8, 'Incorrect reconstruction using V @ diag(e) @ V.T')
            else:
                eigvals, _ = torch.symeig(x, eigenvectors=True, upper=upper)
                self.assertEqual(eigvals, oute, 'Eigenvalues mismatch')
                self.assertEqual(torch.empty(0, device=device, dtype=dtype), outv, 'Eigenvector matrix not empty')

            rese, resv = x.symeig(eigenvectors=eigenvectors, upper=upper)
            self.assertEqual(rese, oute, "outputs of symeig and symeig with out don't match")
            self.assertEqual(resv, outv, "outputs of symeig and symeig with out don't match")

            # test non-contiguous
            x = random_symmetric_matrix(*dims, dtype=dtype, device=device)
            n_dim = len(dims) + 1
            # Reverse the batch dimensions and the matrix dimensions and then concat them
            x = x.permute(tuple(range(n_dim - 3, -1, -1)) + (n_dim - 1, n_dim - 2))
            assert not x.is_contiguous(), "x is intentionally non-contiguous"
            rese, resv = torch.symeig(x, eigenvectors=eigenvectors, upper=upper)
            if eigenvectors:
                x_recon = torch.matmul(torch.matmul(resv, torch.diag_embed(rese)), resv.transpose(-2, -1))
                self.assertEqual(x, x_recon, 1e-8, 'Incorrect reconstruction using V @ diag(e) @ V.T')
            else:
                eigvals, _ = torch.symeig(x, eigenvectors=True, upper=upper)
                self.assertEqual(eigvals, rese, 'Eigenvalues mismatch')
                self.assertEqual(torch.empty(0, device=device, dtype=dtype), resv, 'Eigenvector matrix not empty')

        batch_dims_set = [(), (3,), (3, 5), (5, 3, 5)]
        for batch_dims, eigenvectors, upper in product(batch_dims_set, (True, False), (True, False)):
            run_test((5,) + batch_dims, eigenvectors, upper)
//  [python]
def test_symeig(self, device, dtype):
        from common_utils import random_symmetric_matrix

        def run_test(dims, eigenvectors, upper):
            x = random_symmetric_matrix(*dims, dtype=dtype, device=device)
            oute = torch.empty(dims[1:] + dims[:1], dtype=dtype, device=device)
            outv = torch.empty(dims[1:] + dims[:1] * 2, dtype=dtype, device=device)
            torch.linalg.eigh(x, eigenvectors=eigenvectors, upper=upper, out=(oute, outv))

            if eigenvectors:
                x_recon = torch.matmul(torch.matmul(outv, torch.diag_embed(oute)), outv.transpose(-2, -1))
                self.assertEqual(x, x_recon, 1e-8, 'Incorrect reconstruction using V @ diag(e) @ V.T')
            else:
                eigvals, _ = torch.linalg.eigh(x, eigenvectors=True, upper=upper)
                self.assertEqual(eigvals, oute, 'Eigenvalues mismatch')
                self.assertEqual(torch.empty(0, device=device, dtype=dtype), outv, 'Eigenvector matrix not empty')

            rese, resv = torch.linalg.eigh(x, eigenvectors=eigenvectors, upper=upper)
            self.assertEqual(rese, oute, "outputs of symeig and symeig with out don't match")
            self.assertEqual(resv, outv, "outputs of symeig and symeig with out don't match")

            # test non-contiguous
            x = random_symmetric_matrix(*dims, dtype=dtype, device=device)
            n_dim = len(dims) + 1
            # Reverse the batch dimensions and the matrix dimensions and then concat them
            x = x.permute(tuple(range(n_dim - 3, -1, -1)) + (n_dim - 1, n_dim - 2))
            assert not x.is_contiguous(), "x is intentionally non-contiguous"
            rese, resv = torch.linalg.eigh(x, eigenvectors=eigenvectors, upper=upper)
            if eigenvectors:
                x_recon = torch.matmul(torch.matmul(resv, torch.diag_embed(rese)), resv.transpose(-2, -1))
                self.assertEqual(x, x_recon, 1e-8, 'Incorrect reconstruction using V @ diag(e) @ V.T')
            else:
                eigvals, _ = torch.linalg.eigh(x, eigenvectors=True, upper=upper)
                self.assertEqual(eigvals, rese, 'Eigenvalues mismatch')
                self.assertEqual(torch.empty(0, device=device, dtype=dtype), resv, 'Eigenvector matrix not empty')

        batch_dims_set = [(), (3,), (3, 5), (5, 3, 5)]
        for batch_dims, eigenvectors, upper in product(batch_dims_set, (True, False), (True, False)):
            run_test((5,) + batch_dims, eigenvectors, upper)
//  [python]
def joint_entropy(x, y, s_x, s_y):
    alpha = 1.01
    # print("x shape is {}, y shape is {}".format(x.shape, y.shape))
    x = calculate_gram_mat(x, s_x)
    y = calculate_gram_mat(y, s_y)
    k = torch.mul(x, y)
    k = k / torch.trace(k)
    print("pairwise distance joint entropy")
    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
    # print("x shape is {}, y shape is {}".format(x.shape, y.shape))
    x = calculate_gram_mat(x, s_x)
    y = calculate_gram_mat(y, s_y)
    k = torch.mul(x, y)
    k = k / torch.trace(k)
    print("pairwise distance joint entropy")
    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 logm_eig(A, spd=True):
    """Batched matrix logarithm through eigenvalue decomposition.

    Parameters
    ----------
    A : torch.Tensor
        Square matrices of shape (B, F, C, T).

    Returns
    -------
    torch.Tensor :
        Matrix logarithm of A.
    """
    e, v = torch.symeig(A, eigenvectors=True)
    e = torch.clamp(e, min=1e-10)  # clamp the eigenvalues to avoid -inf
    return v @ torch.diag_embed(
        torch.log(e), dim1=2, dim2=3) @ v.transpose(2, 3)
//  [python]
def logm_eig(A, spd=True):
    """Batched matrix logarithm through eigenvalue decomposition.

    Parameters
    ----------
    A : torch.Tensor
        Square matrices of shape (B, F, C, T).

    Returns
    -------
    torch.Tensor :
        Matrix logarithm of A.
    """
    e, v = torch.linalg.eigh(A)
    e = torch.clamp(e, min=1e-10)  # clamp the eigenvalues to avoid -inf
    return v @ torch.diag_embed(
        torch.log(e), dim1=2, dim2=3) @ v.transpose(2, 3)
//  [python]
def torch_mat2quat(M):
    all_pose = []
    for i in range(M.size()[0]):
        rotate = M[i, :3, :3]
        translate = M[i, :3, 3]

        # Qyx refers to the contribution of the y input vector component to
        # the x output vector component.  Qyx is therefore the same as
        # M[0,1].  The notation is from the Wikipedia article.
        Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz = rotate.flatten()
        #     print(Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz)
        # Fill only lower half of symmetric matrix
        K = torch.tensor([[Qxx - Qyy - Qzz, 0, 0, 0], [Qyx + Qxy, Qyy - Qxx - Qzz, 0, 0], [Qzx + Qxz, Qzy + Qyz, Qzz - Qxx - Qyy, 0],
                          [Qyz - Qzy, Qzx - Qxz, Qxy - Qyx, Qxx + Qyy + Qzz]]) / 3.0
        K = K.to(M.device)
        # Use Hermitian eigenvectors, values for speed
        vals, vecs = torch.symeig(K, True, False)
        # Select largest eigenvector, reorder to w,x,y,z quaternion

        q = vecs[[3, 0, 1, 2], torch.argmax(vals)]
        # Prefer quaternion with positive w
        # (q * -1 corresponds to same rotation as q)
        if q[0] < 0:
            q *= -1

        pose = torch.cat((q, translate), dim=0)
        all_pose.append(pose)
    all_pose = torch.stack(all_pose, dim=0)
    return all_pose
//  [python]
def torch_mat2quat(M):
    all_pose = []
    for i in range(M.size()[0]):
        rotate = M[i, :3, :3]
        translate = M[i, :3, 3]

        # Qyx refers to the contribution of the y input vector component to
        # the x output vector component.  Qyx is therefore the same as
        # M[0,1].  The notation is from the Wikipedia article.
        Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz = rotate.flatten()
        #     print(Qxx, Qyx, Qzx, Qxy, Qyy, Qzy, Qxz, Qyz, Qzz)
        # Fill only lower half of symmetric matrix
        K = torch.tensor([[Qxx - Qyy - Qzz, 0, 0, 0], [Qyx + Qxy, Qyy - Qxx - Qzz, 0, 0], [Qzx + Qxz, Qzy + Qyz, Qzz - Qxx - Qyy, 0],
                          [Qyz - Qzy, Qzx - Qxz, Qxy - Qyx, Qxx + Qyy + Qzz]]) / 3.0
        K = K.to(M.device)
        # Use Hermitian eigenvectors, values for speed
        vals, vecs = torch.linalg.eigh(K)
        # Select largest eigenvector, reorder to w,x,y,z quaternion

        q = vecs[[3, 0, 1, 2], torch.argmax(vals)]
        # Prefer quaternion with positive w
        # (q * -1 corresponds to same rotation as q)
        if q[0] < 0:
            q *= -1

        pose = torch.cat((q, translate), dim=0)
        all_pose.append(pose)
    all_pose = torch.stack(all_pose, dim=0)
    return all_pose
//  [python]
def _get_style_vals(self, tensor):
        mean, cov = self._calc_2_moments(tensor) 
        if mean is None:
            return None, None, None
        eigvals, eigvects = torch.symeig(cov, eigenvectors=True)
        eigroot_mat = torch.diag(torch.sqrt(eigvals.clamp(min=0)))     
        root_cov = torch.mm(torch.mm(eigvects, eigroot_mat), eigvects.t())  
        tr_cov = eigvals.clamp(min=0).sum() 
        return mean, tr_cov, root_cov
//  [python]
def _get_style_vals(self, tensor):
        mean, cov = self._calc_2_moments(tensor) 
        if mean is None:
            return None, None, None
        eigvals, eigvects = torch.linalg.eigh(cov)
        eigroot_mat = torch.diag(torch.sqrt(eigvals.clamp(min=0)))     
        root_cov = torch.mm(torch.mm(eigvects, eigroot_mat), eigvects.t())  
        tr_cov = eigvals.clamp(min=0).sum() 
        return mean, tr_cov, root_cov
//  [python]
def batch_symeig(X):
    # it is much faster to run symeig on CPU
    X = X.cpu()
    b, d, _ = X.size()
    bv = X.new(b,d,d)
   
    for batch_idx in range(X.shape[0]):
        e,v = torch.symeig(X[batch_idx,:,:].squeeze(), True)
        bv[batch_idx,:,:] = v
    bv = bv.cuda()
    return bv
//  [python]
def batch_symeig(X):
    # it is much faster to run symeig on CPU
    X = X.cpu()
    b, d, _ = X.size()
    bv = X.new(b,d,d)
   
    for batch_idx in range(X.shape[0]):
        e,v = torch.linalg.eigh(X[batch_idx,:,:].squeeze())
        bv[batch_idx,:,:] = v
    bv = bv.cuda()
    return bv
//  [python]
def _calc_l2wass_dist(self, mean_stl, tr_cov_stl, root_cov_stl, mean_synth, cov_synth):
        tr_cov_synth = torch.symeig(cov_synth, eigenvectors=True)[0].clamp(min=0).sum()
        mean_diff_squared = (mean_stl - mean_synth).pow(2).sum()
        cov_prod = torch.mm(torch.mm(root_cov_stl, cov_synth), root_cov_stl)
        var_overlap = torch.sqrt(torch.symeig(cov_prod, eigenvectors=True)[0].clamp(min=0)+1e-8).sum()
        dist = mean_diff_squared + tr_cov_stl + tr_cov_synth - 2*var_overlap
        return dist
//  [python]
def _calc_l2wass_dist(self, mean_stl, tr_cov_stl, root_cov_stl, mean_synth, cov_synth):
        tr_cov_synth = torch.linalg.eigh(cov_synth)[0].clamp(min=0).sum()
        mean_diff_squared = (mean_stl - mean_synth).pow(2).sum()
        cov_prod = torch.mm(torch.mm(root_cov_stl, cov_synth), root_cov_stl)
        var_overlap = torch.sqrt(torch.linalg.eigh(cov_prod)[0].clamp(min=0)+1e-8).sum()
        dist = mean_diff_squared + tr_cov_stl + tr_cov_synth - 2*var_overlap
        return dist
//  [python]
def main(device):
    x = torch.rand(4, requires_grad=True, device=device)
    x1 = x[0].pow(2) + x[1].pow(2)
    x2 = x[2].pow(2) + x[3].pow(2)
    y = x1 * x2

    h = hessian(y, x)
    print(h)

    h_exact = torch.tensor([
        [2 * x2, 0, 4 * x[0] * x[2], 4 * x[0] * x[3]],
        [0, 2 * x2, 4 * x[1] * x[2], 4 * x[1] * x[3]],
        [4 * x[0] * x[2], 4 * x[1] * x[2], 2 * x1, 0],
        [4 * x[0] * x[3], 4 * x[1] * x[3], 0, 2 * x1],
    ])
    print(h_exact)

    e, v = torch.symeig(h, eigenvectors=True)
    print(e)
    print(v)
//  [python]
def main(device):
    x = torch.rand(4, requires_grad=True, device=device)
    x1 = x[0].pow(2) + x[1].pow(2)
    x2 = x[2].pow(2) + x[3].pow(2)
    y = x1 * x2

    h = hessian(y, x)
    print(h)

    h_exact = torch.tensor([
        [2 * x2, 0, 4 * x[0] * x[2], 4 * x[0] * x[3]],
        [0, 2 * x2, 4 * x[1] * x[2], 4 * x[1] * x[3]],
        [4 * x[0] * x[2], 4 * x[1] * x[2], 2 * x1, 0],
        [4 * x[0] * x[3], 4 * x[1] * x[3], 0, 2 * x1],
    ])
    print(h_exact)

    e, v = torch.linalg.eigh(h)
    print(e)
    print(v)
//  [python]
def compute_eig_vectors_only(hidden, only_last=False):

    if only_last:
        hTemp = hidden[-1] - torch.mean(hidden[-1], 0)
        cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
        cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
        cov = (cov + cov.transpose(1, 0)) / 2
        _, vecTemp = torch.symeig(cov, eigenvectors=True)
        eigVec = vecTemp.float()
    else:
        eigVec = []
        for idx in range(len(hidden)):
            hTemp = hidden[idx] - torch.mean(hidden[idx], 0)
            cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
            cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
            cov = (cov + cov.transpose(1, 0)) / 2
            # _, _, vecTemp = torch.svd(cov, compute_uv=True)  # compute eigenvectors and values
            _, vecTemp = torch.symeig(cov, eigenvectors=True)
            vecTemp = vecTemp.float()
            eigVec.append(vecTemp.cpu())

    return eigVec
//  [python]
def compute_eig_vectors_only(hidden, only_last=False):

    if only_last:
        hTemp = hidden[-1] - torch.mean(hidden[-1], 0)
        cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
        cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
        cov = (cov + cov.transpose(1, 0)) / 2
        _, vecTemp = torch.linalg.eigh(cov)
        eigVec = vecTemp.float()
    else:
        eigVec = []
        for idx in range(len(hidden)):
            hTemp = hidden[idx] - torch.mean(hidden[idx], 0)
            cov = hTemp.transpose(1, 0) @ hTemp / hTemp.shape[0]  # compute covariance matrix
            cov = cov.double()  # cast as 64bit to mitigate underflow in matrix factorization
            cov = (cov + cov.transpose(1, 0)) / 2
            # _, _, vecTemp = torch.svd(cov, compute_uv=True)  # compute eigenvectors and values
            _, vecTemp = torch.linalg.eigh(cov)
            vecTemp = vecTemp.float()
            eigVec.append(vecTemp.cpu())

    return eigVec
//  [python]
def _generalized_eigenvalue_decomposition(a, b):
    cholesky = stable_cholesky(b)
    inv_cholesky = torch.inverse(cholesky)
    cmat = inv_cholesky @ a @ inv_cholesky.conj().transpose(-1, -2)
    e_val, e_vec = torch.symeig(cmat, eigenvectors=True)
    e_vec = torch.matmul(inv_cholesky.conj().transpose(-1, -2), e_vec)
    return e_val, e_vec
//  [python]
def _generalized_eigenvalue_decomposition(a, b):
    cholesky = stable_cholesky(b)
    inv_cholesky = torch.inverse(cholesky)
    cmat = inv_cholesky @ a @ inv_cholesky.conj().transpose(-1, -2)
    e_val, e_vec = torch.linalg.eigh(cmat)
    e_vec = torch.matmul(inv_cholesky.conj().transpose(-1, -2), e_vec)
    return e_val, e_vec
//  [python]
def forward(self, mix: torch.Tensor, target_scm: torch.Tensor, noise_scm: torch.Tensor):
        """Compute and apply MVDR beamformer from the speech and noise SCM matrices.

        :math:`\\mathbf{w} =  \\displaystyle \\frac{\\Sigma_{nn}^{-1} \\mathbf{a}}{
        \\mathbf{a}^H \\Sigma_{nn}^{-1} \\mathbf{a}}` where :math:`\\mathbf{a}` is the
        ATF estimated from the target SCM.

        Args:
            mix (torch.ComplexTensor): shape (batch, mics, freqs, frames)
            target_scm (torch.ComplexTensor): (batch, mics, mics, freqs)
            noise_scm (torch.ComplexTensor): (batch, mics, mics, freqs)

        Returns:
            Filtered mixture. torch.ComplexTensor (batch, freqs, frames)
        """
        e_val, e_vec = torch.symeig(target_scm.permute(0, 3, 1, 2), eigenvectors=True)
        rtf_vect = e_vec[..., -1]
        return self.from_rtf_vect(mix=mix, rtf_vec=rtf_vect.transpose(-1, -2), noise_scm=noise_scm)
//  [python]
def forward(self, mix: torch.Tensor, target_scm: torch.Tensor, noise_scm: torch.Tensor):
        """Compute and apply MVDR beamformer from the speech and noise SCM matrices.

        :math:`\\mathbf{w} =  \\displaystyle \\frac{\\Sigma_{nn}^{-1} \\mathbf{a}}{
        \\mathbf{a}^H \\Sigma_{nn}^{-1} \\mathbf{a}}` where :math:`\\mathbf{a}` is the
        ATF estimated from the target SCM.

        Args:
            mix (torch.ComplexTensor): shape (batch, mics, freqs, frames)
            target_scm (torch.ComplexTensor): (batch, mics, mics, freqs)
            noise_scm (torch.ComplexTensor): (batch, mics, mics, freqs)

        Returns:
            Filtered mixture. torch.ComplexTensor (batch, freqs, frames)
        """
        e_val, e_vec = torch.linalg.eigh(target_scm.permute(0, 3, 1, 2))
        rtf_vect = e_vec[..., -1]
        return self.from_rtf_vect(mix=mix, rtf_vec=rtf_vect.transpose(-1, -2), noise_scm=noise_scm)
//  [python]
def train_cov(X, model, optimizer, maxiter, noise_std):
    model.initialize()
    weights = []
    print('======GLD Leading======')
    for i in range(maxiter):
        covariance = eval_Cov(X, model)
        covariance = covariance / covariance.trace() * noise_std
        e, v = torch.symeig(covariance, eigenvectors=True)

        loss = model(X).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        model.weight.data += torch.matmul(v, torch.randn(model.n_weights).to(model.weight.device) * e.sqrt())
        weights.append(model.weight.data.cpu().clone().numpy())
        if i % 100 == 0:
            print('Step:', i, 'Loss:', loss.item(), 'Location:', model.weight.data)
    return np.stack(weights)
//  [python]
def train_cov(X, model, optimizer, maxiter, noise_std):
    model.initialize()
    weights = []
    print('======GLD Leading======')
    for i in range(maxiter):
        covariance = eval_Cov(X, model)
        covariance = covariance / covariance.trace() * noise_std
        e, v = torch.linalg.eigh(covariance)

        loss = model(X).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        model.weight.data += torch.matmul(v, torch.randn(model.n_weights).to(model.weight.device) * e.sqrt())
        weights.append(model.weight.data.cpu().clone().numpy())
        if i % 100 == 0:
            print('Step:', i, 'Loss:', loss.item(), 'Location:', model.weight.data)
    return np.stack(weights)
//  [python]
def train_hess(X, model, optimizer, maxiter, noise_std):
    model.initialize()
    weights = []
    print('======GLD Hess======')
    for i in range(maxiter):
        hessian = eval_Hess(X, model)
        hessian = hessian / hessian.trace() * noise_std
        e, v = torch.symeig(hessian, eigenvectors=True)

        loss = model(X).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        model.weight.data += torch.matmul(v, torch.randn(model.n_weights).to(model.weight.device) * e.sqrt())
        weights.append(model.weight.data.cpu().clone().numpy())
        if i % 100 == 0:
            print('Step:', i, 'Loss:', loss.item(), 'Location:', model.weight.data)
    return np.stack(weights)
//  [python]
def train_hess(X, model, optimizer, maxiter, noise_std):
    model.initialize()
    weights = []
    print('======GLD Hess======')
    for i in range(maxiter):
        hessian = eval_Hess(X, model)
        hessian = hessian / hessian.trace() * noise_std
        e, v = torch.linalg.eigh(hessian)

        loss = model(X).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        model.weight.data += torch.matmul(v, torch.randn(model.n_weights).to(model.weight.device) * e.sqrt())
        weights.append(model.weight.data.cpu().clone().numpy())
        if i % 100 == 0:
            print('Step:', i, 'Loss:', loss.item(), 'Location:', model.weight.data)
    return np.stack(weights)
//  [python]
def train_hessMax(X, model, optimizer, maxiter, noise_std):
    model.initialize()
    weights = []
    print('======GLD HessMax======')
    for i in range(maxiter):
        hessian = eval_Hess(X, model)
        e, v = torch.symeig(hessian, eigenvectors=True)
        std = math.sqrt(noise_std)

        loss = model(X).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        model.weight.data += v[:,-1] * torch.randn(1).to(model.weight.device) * std
        weights.append(model.weight.data.cpu().clone().numpy())
        if i % 100 == 0:
            print('Step:', i, 'Loss:', loss.item(), 'Location:', model.weight.data[0], model.weight.data[1],
                    'v:', v[0,0], v[0,1], v[1,0], v[1,1], 'e:', e[0], e[1])
    return np.stack(weights)
//  [python]
def train_hessMax(X, model, optimizer, maxiter, noise_std):
    model.initialize()
    weights = []
    print('======GLD HessMax======')
    for i in range(maxiter):
        hessian = eval_Hess(X, model)
        e, v = torch.linalg.eigh(hessian)
        std = math.sqrt(noise_std)

        loss = model(X).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        model.weight.data += v[:,-1] * torch.randn(1).to(model.weight.device) * std
        weights.append(model.weight.data.cpu().clone().numpy())
        if i % 100 == 0:
            print('Step:', i, 'Loss:', loss.item(), 'Location:', model.weight.data[0], model.weight.data[1],
                    'v:', v[0,0], v[0,1], v[1,0], v[1,1], 'e:', e[0], e[1])
    return np.stack(weights)
//  [python]
def train(self, x):
        '''training pca.
        Args:
            x: [N, dim] FloatTensor containing data which undergoes PCA/Whitening.
        '''

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

        # x = x.double()
        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.n_components < x2.size(0)):
            k_indices = torch.argsort(L, descending=True)[:self.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('===============================================')

        # save features, labels to h5 file.
        filename = os.path.join(self.parameters_path)
        np.savez(filename, U=U, lams=lams, mu=mu, Utmu=Utmu)
//  [python]
def train(self, x):
        '''training pca.
        Args:
            x: [N, dim] FloatTensor containing data which undergoes PCA/Whitening.
        '''

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

        # x = x.double()
        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.n_components < x2.size(0)):
            k_indices = torch.argsort(L, descending=True)[:self.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('===============================================')

        # save features, labels to h5 file.
        filename = os.path.join(self.parameters_path)
        np.savez(filename, U=U, lams=lams, mu=mu, Utmu=Utmu)
//  [python]
def average_quaternions(quaternions: torch.Tensor,
                        weights: Optional[torch.Tensor] = None,
                        keepdim: bool = False,
                        ):
    """
    Calculate average quaternion with the method described in the paper
    Markley, F Landis, Yang Cheng, John L Crassidis, and Yaakov Oshman. “Averaging Quaternions.”
    https://doi.org/10.2514/1.28949.

    Args:
        quaternions: ... x N x 4 - tensor and contains the quaternions to average in the rows.
            The quaternions are arranged as (x,y,z,w), with w being the scalar
        weights: ... x N - weight of each quaternion
        keepdim:

    Returns:
            torch.Tensor: ... x 4
    """

    # Number of quaternions to average
    samples = quaternions.shape[-2]
    batch_shape = quaternions.shape[:-2]
    quaternions = quaternions.reshape(-1, samples, 4)

    if weights is None:
        weights = 1. / samples
    else:
        assert weights.shape[-1] == samples, \
            f'number of weights {weights.shape[-1]} does not equal to number of quaternions {samples}'
        weights = weights / weights.sum(dim=-1, keepdim=True).clamp(min=1e-7)
        weights = weights.reshape(-1, samples, 1).to(quaternions.device)

    # B x 4 X 4
    mat_a = torch.einsum('bni,bnj->bij', [quaternions * weights, quaternions])
    # compute eigenvalues and -vectors
    _, eigen_vectors = torch.symeig(mat_a, eigenvectors=True)
    max_value_eigenvector = eigen_vectors[..., -1].reshape(*batch_shape, 4)

    return max_value_eigenvector.unsqueeze(-2) if keepdim else max_value_eigenvector