jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, generator, data=None):
        self.generator = generator
        if data is None:
            evecs = dict()
            diags = dict()

            kfac_blocks = generator.get_kfac_blocks()
            for layer_id, layer in \
                    self.generator.layer_collection.layers.items():
                a, g = kfac_blocks[layer_id]
                evals_a, evecs_a = torch.symeig(a, eigenvectors=True)
                evals_g, evecs_g = torch.symeig(g, eigenvectors=True)
                evecs[layer_id] = (evecs_a, evecs_g)
                diags[layer_id] = kronecker(evals_g.view(-1, 1),
                                            evals_a.view(-1, 1))
                del a, g, kfac_blocks[layer_id]
            self.data = (evecs, diags)
        else:
            self.data = data
//  [python]
def __init__(self, generator, data=None):
        self.generator = generator
        if data is None:
            evecs = dict()
            diags = dict()

            kfac_blocks = generator.get_kfac_blocks()
            for layer_id, layer in \
                    self.generator.layer_collection.layers.items():
                a, g = kfac_blocks[layer_id]
                evals_a, evecs_a = torch.linalg.eigh(a)
                evals_g, evecs_g = torch.linalg.eigh(g)
                evecs[layer_id] = (evecs_a, evecs_g)
                diags[layer_id] = kronecker(evals_g.view(-1, 1),
                                            evals_a.view(-1, 1))
                del a, g, kfac_blocks[layer_id]
            self.data = (evecs, diags)
        else:
            self.data = data
//  [python]
def __init__(self, n_variables, highest_order, random_seed_pair=(None, None)):
		case_seed, init_seed = random_seed_pair
		self.suggested_init = torch.empty(0).long()
		self.suggested_init = torch.cat([self.suggested_init, sample_init_points([2] * n_variables, 20 - self.suggested_init.size(0), init_seed).long()], dim=0)
		self.n_variables = n_variables
		self.highest_order = highest_order
		self.interaction_coef = generate_function_on_highorderbinary(n_variables=n_variables, highest_order=highest_order, random_seed=case_seed)
		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(self.n_variables):
			adjmat = torch.diag(torch.ones(1), -1) + torch.diag(torch.ones(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, n_variables, highest_order, random_seed_pair=(None, None)):
		case_seed, init_seed = random_seed_pair
		self.suggested_init = torch.empty(0).long()
		self.suggested_init = torch.cat([self.suggested_init, sample_init_points([2] * n_variables, 20 - self.suggested_init.size(0), init_seed).long()], dim=0)
		self.n_variables = n_variables
		self.highest_order = highest_order
		self.interaction_coef = generate_function_on_highorderbinary(n_variables=n_variables, highest_order=highest_order, random_seed=case_seed)
		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(self.n_variables):
			adjmat = torch.diag(torch.ones(1), -1) + torch.diag(torch.ones(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 compute_eigendecomposition(self, impl='symeig'):
        if impl == 'symeig':
            self.evals, V = torch.symeig(torch.mm(self.data, self.data.t()),
                                         eigenvectors=True)
            self.evecs = torch.mm(self.data.t(), V) / \
                (self.evals**.5).unsqueeze(0)
        else:
            raise NotImplementedError
//  [python]
def compute_eigendecomposition(self, impl='symeig'):
        if impl == 'symeig':
            self.evals, V = torch.linalg.eigh(torch.mm(self.data, self.data.t()))
            self.evecs = torch.mm(self.data.t(), V) / \
                (self.evals**.5).unsqueeze(0)
        else:
            raise NotImplementedError
//  [python]
def test_symeig(self):
        xval = torch.rand(100, 3)
        cov = torch.mm(xval.t(), xval)
        rese = torch.zeros(3)
        resv = torch.zeros(3, 3)

        # First call to symeig
        self.assertTrue(resv.is_contiguous(), 'resv is not contiguous')
        torch.symeig(cov.clone(), True, out=(rese, resv))
        ahat = torch.mm(torch.mm(resv, torch.diag(rese)), resv.t())
        self.assertEqual(cov, ahat, 1e-8, 'VeV\' wrong')

        # Second call to symeig
        self.assertFalse(resv.is_contiguous(), 'resv is contiguous')
        torch.symeig(cov.clone(), True, out=(rese, resv))
        ahat = torch.mm(torch.mm(resv, torch.diag(rese)), resv.t())
        self.assertEqual(cov, ahat, 1e-8, 'VeV\' wrong')

        # test non-contiguous
        X = torch.rand(5, 5)
        X = X.t() * X
        e = torch.zeros(4, 2).select(1, 1)
        v = torch.zeros(4, 2, 4)[:, 1]
        self.assertFalse(v.is_contiguous(), 'V is contiguous')
        self.assertFalse(e.is_contiguous(), 'E is contiguous')
        torch.symeig(X, True, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e)), v.t())
        self.assertEqual(X, Xhat, 1e-8, 'VeV\' wrong')
//  [python]
def test_symeig(self):
        xval = torch.rand(100, 3)
        cov = torch.mm(xval.t(), xval)
        rese = torch.zeros(3)
        resv = torch.zeros(3, 3)

        # First call to symeig
        self.assertTrue(resv.is_contiguous(), 'resv is not contiguous')
        e, v = torch.linalg.eigh(cov.clone())
        rese = e
        resv = v
        ahat = torch.mm(torch.mm(resv, torch.diag(rese)), resv.t())
        self.assertEqual(cov, ahat, 1e-8, 'VeV\' wrong')

        # Second call to symeig
        self.assertFalse(resv.is_contiguous(), 'resv is contiguous')
        e, v = torch.linalg.eigh(cov.clone())
        rese = e
        resv = v
        ahat = torch.mm(torch.mm(resv, torch.diag(rese)), resv.t())
        self.assertEqual(cov, ahat, 1e-8, 'VeV\' wrong')

        # test non-contiguous
        X = torch.rand(5, 5)
        X = X.t() * X
        e = torch.zeros(4, 2).select(1, 1)
        v = torch.zeros(4, 2, 4)[:, 1]
        self.assertFalse(v.is_contiguous(), 'V is contiguous')
        self.assertFalse(e.is_contiguous(), 'E is contiguous')
        e, v = torch.linalg.eigh(X)
        Xhat = torch.mm(torch.mm(v, torch.diag(e)), v.t())
        self.assertEqual(X, Xhat, 1e-8, 'VeV\' wrong')
//  [python]
def pca_eigen_values(x: torch.Tensor):
    temp = x[:, :3] - x.mean(dim=0)[None, :3]
    cov = (temp.transpose(0, 1) @ temp) / x.shape[0]
    e, v = torch.symeig(cov, eigenvectors=True)
    n = v[:, 0]
    return e[0:1], n
//  [python]
def pca_eigen_values(x: torch.Tensor):
    temp = x[:, :3] - x.mean(dim=0)[None, :3]
    cov = (temp.transpose(0, 1) @ temp) / x.shape[0]
    e, v = torch.linalg.eigh(cov)
    n = v[:, 0]
    return e[0:1], n
//  [python]
def _update_T(self, m):
        """Do eigen decomposition for computing inverse of the ~ fisher.
        :param m: The layer
        :return: no returns.
        """
        if self.svd:
           d_x, Q_x = torch.symeig(self.m_xx[m], eigenvectors=True)
           dampening=d_x.max()*self.dampening           
           d_x=1/(d_x+dampening)**(0.5) 
           self.T[m]= ( Q_x*(d_x.unsqueeze(0))) @ Q_x.t()           
        else:
           max_ev_xx=max_eignvalue(self.m_xx[m])
           dampening=max_ev_xx**self.dampening 
           self.T[m]=isqrt_newton_schulz(self.m_xx[m]+dampening*torch.eye(self.m_xx[m].size(0)).to(self.m_xx[m]))
//  [python]
def _update_T(self, m):
        """Do eigen decomposition for computing inverse of the ~ fisher.
        :param m: The layer
        :return: no returns.
        """
        if self.svd:
           d_x, Q_x = torch.linalg.eigh(self.m_xx[m])
           dampening=d_x.max()*self.dampening           
           d_x=1/(d_x+dampening)**(0.5) 
           self.T[m]= ( Q_x*(d_x.unsqueeze(0))) @ Q_x.t()           
        else:
           max_ev_xx=max_eignvalue(self.m_xx[m])
           dampening=max_ev_xx**self.dampening 
           self.T[m]=isqrt_newton_schulz(self.m_xx[m]+dampening*torch.eye(self.m_xx[m].size(0)).to(self.m_xx[m]))
//  [python]
def test_symeig_out_errors_and_warnings(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        # if non-empty out tensor with wrong shape is passed a warning is given
        a = random_hermitian_matrix(3, dtype=dtype, device=device)
        real_dtype = a.real.dtype if dtype.is_complex else dtype
        out_w = torch.empty(7, 7, dtype=real_dtype, device=device)
        out_v = torch.empty(7, 7, dtype=dtype, device=device)
        with warnings.catch_warnings(record=True) as w:
            # Trigger warning
            torch.symeig(a, out=(out_w, out_v))
            self.assertTrue("An output with one or more elements was resized" in str(w[-2].message))
            self.assertTrue("An output with one or more elements was resized" in str(w[-1].message))

        # dtypes should be safely castable
        out_w = torch.empty(0, dtype=real_dtype, device=device)
        out_v = torch.empty(0, dtype=torch.int, device=device)
        with self.assertRaisesRegex(RuntimeError, "but got eigenvectors with dtype Int"):
            torch.symeig(a, out=(out_w, out_v))

        out_w = torch.empty(0, dtype=torch.int, device=device)
        out_v = torch.empty(0, dtype=dtype, device=device)
        with self.assertRaisesRegex(RuntimeError, "but got eigenvalues with dtype Int"):
            torch.symeig(a, out=(out_w, out_v))

        # device should match
        if torch.cuda.is_available():
            wrong_device = 'cpu' if self.device_type != 'cpu' else 'cuda'
            out_w = torch.empty(0, device=wrong_device, dtype=dtype)
            out_v = torch.empty(0, device=device, dtype=dtype)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.symeig(a, out=(out_w, out_v))
            out_w = torch.empty(0, device=device, dtype=dtype)
            out_v = torch.empty(0, device=wrong_device, dtype=dtype)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.symeig(a, out=(out_w, out_v))
//  [python]
def test_symeig_out_errors_and_warnings(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        # if non-empty out tensor with wrong shape is passed a warning is given
        a = random_hermitian_matrix(3, dtype=dtype, device=device)
        real_dtype = a.real.dtype if dtype.is_complex else dtype
        out_w = torch.empty(7, 7, dtype=real_dtype, device=device)
        out_v = torch.empty(7, 7, dtype=dtype, device=device)
        with warnings.catch_warnings(record=True) as w:
            # Trigger warning
            torch.linalg.eigh(a, out=(out_w, out_v))
            self.assertTrue("An output with one or more elements was resized" in str(w[-2].message))
            self.assertTrue("An output with one or more elements was resized" in str(w[-1].message))

        # dtypes should be safely castable
        out_w = torch.empty(0, dtype=real_dtype, device=device)
        out_v = torch.empty(0, dtype=torch.int, device=device)
        with self.assertRaisesRegex(RuntimeError, "but got eigenvectors with dtype Int"):
            torch.linalg.eigh(a, out=(out_w, out_v))

        out_w = torch.empty(0, dtype=torch.int, device=device)
        out_v = torch.empty(0, dtype=dtype, device=device)
        with self.assertRaisesRegex(RuntimeError, "but got eigenvalues with dtype Int"):
            torch.linalg.eigh(a, out=(out_w, out_v))

        # device should match
        if torch.cuda.is_available():
            wrong_device = 'cpu' if self.device_type != 'cpu' else 'cuda'
            out_w = torch.empty(0, device=wrong_device, dtype=dtype)
            out_v = torch.empty(0, device=device, dtype=dtype)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.linalg.eigh(a, out=(out_w, out_v))
            out_w = torch.empty(0, device=device, dtype=dtype)
            out_v = torch.empty(0, device=wrong_device, dtype=dtype)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.linalg.eigh(a, out=(out_w, out_v))
//  [python]
def test_symeig(self, device):
        def op(x):
            return torch.symeig(x, eigenvectors=True)[0]

        x = torch.randn(3, 3, device=device, requires_grad=True)
        self._batched_grad_test(op, (x,), {})
        self._batched_grad_grad_test(op, (x,), {})
//  [python]
def test_symeig(self, device):
        def op(x):
            return torch.linalg.eigh(x, eigenvectors=True)[0]

        x = torch.randn(3, 3, device=device, requires_grad=True)
        self._batched_grad_test(op, (x,), {})
        self._batched_grad_grad_test(op, (x,), {})
//  [python]
def rotate_to_principle_components(x: torch.Tensor, scale=True):
    temp = x[:, :3] - x.mean(dim=0)[None, :3]
    cov = temp.transpose(0, 1) @ temp / x.shape[0]
    e, v = torch.symeig(cov, eigenvectors=True)

    # rotate xyz
    rotated = x[:, :3]@v
    if scale:
        # scale to unit var on for the larger eigen value
        rotated = rotated / torch.sqrt(e[2])

    # if x contains normals rotate the normals as well
    if x.shape[1] == 6:
        rotated = torch.cat([rotated, x[:, 3:]@v], dim=-1)
    return rotated
//  [python]
def rotate_to_principle_components(x: torch.Tensor, scale=True):
    temp = x[:, :3] - x.mean(dim=0)[None, :3]
    cov = temp.transpose(0, 1) @ temp / x.shape[0]
    e, v = torch.linalg.eigh(cov)

    # rotate xyz
    rotated = x[:, :3]@v
    if scale:
        # scale to unit var on for the larger eigen value
        rotated = rotated / torch.sqrt(e[2])

    # if x contains normals rotate the normals as well
    if x.shape[1] == 6:
        rotated = torch.cat([rotated, x[:, 3:]@v], dim=-1)
    return rotated
//  [python]
def estimate_normals_torch(inputpc, max_nn):
    from torch_cluster import knn_graph
    knn = knn_graph(inputpc[:, :3], max_nn, loop=False)
    knn = knn.view(2, inputpc.shape[0], max_nn)[0]
    x = inputpc[knn][:, :, :3]
    temp = x[:, :, :3] - x.mean(dim=1)[:, None, :3]
    cov = temp.transpose(1, 2) @ temp / x.shape[0]
    e, v = torch.symeig(cov, eigenvectors=True)
    n = v[:, :, 0]
    return torch.cat([inputpc[:, :3], n], dim=-1)
//  [python]
def estimate_normals_torch(inputpc, max_nn):
    from torch_cluster import knn_graph
    knn = knn_graph(inputpc[:, :3], max_nn, loop=False)
    knn = knn.view(2, inputpc.shape[0], max_nn)[0]
    x = inputpc[knn][:, :, :3]
    temp = x[:, :, :3] - x.mean(dim=1)[:, None, :3]
    cov = temp.transpose(1, 2) @ temp / x.shape[0]
    e, v = torch.linalg.eigh(cov)
    n = v[:, :, 0]
    return torch.cat([inputpc[:, :3], n], dim=-1)
//  [python]
def mds(D, d):
    """
    Args:
    D - (..., n, n) distance matrix

    Returns:
    X - (..., n, d) hyperbolic embeddings
    """
    Y = -torch.cosh(D)
    # print("Y:", Y)
    eigenvals, eigenvecs = torch.symeig(Y, eigenvectors=True)
    # print(Y.shape, eigenvals.shape, eigenvecs.shape)
    # print(eigenvals, eigenvecs)
    X = torch.sqrt(torch.clamp(eigenvals[-d:], min=0.)) * eigenvecs[..., -d:]
    # print("testing")
    # print(X)
    # print(Y @ X)
    u = torch.sqrt(1 + torch.sum(X * X, dim=-1, keepdim=True))
    M = torch.cat((u, X), dim=-1)
    # print(minkowski.pairwise_bilinear_pairing(M, M))
    return torch.cat((u, X), dim=-1)
//  [python]
def mds(D, d):
    """
    Args:
    D - (..., n, n) distance matrix

    Returns:
    X - (..., n, d) hyperbolic embeddings
    """
    Y = -torch.cosh(D)
    eigenvals, eigenvecs = torch.linalg.eigh(Y)
    X = torch.sqrt(torch.clamp(eigenvals[..., -d:], min=0.)) * eigenvecs[..., -d:]
    u = torch.sqrt(1 + torch.sum(X * X, dim=-1, keepdim=True))
    M = torch.cat((u, X), dim=-1)
    return torch.cat((u, X), dim=-1)
//  [python]
def __init__(self, generator, data=None, examples=None):
        self._check_data_examples(data, examples)

        self.generator = generator
        if data is None:
            evecs = dict()
            diags = dict()

            kfac_blocks = generator.get_kfac_blocks(examples)
            for layer_id, layer in \
                    self.generator.layer_collection.layers.items():
                a, g = kfac_blocks[layer_id]
                evals_a, evecs_a = torch.symeig(a, eigenvectors=True)
                evals_g, evecs_g = torch.symeig(g, eigenvectors=True)
                evecs[layer_id] = (evecs_a, evecs_g)
                diags[layer_id] = kronecker(evals_g.view(-1, 1),
                                            evals_a.view(-1, 1))
                del a, g, kfac_blocks[layer_id]
            self.data = (evecs, diags)
        else:
            self.data = data
//  [python]
def __init__(self, generator, data=None, examples=None):
        self._check_data_examples(data, examples)

        self.generator = generator
        if data is None:
            evecs = dict()
            diags = dict()

            kfac_blocks = generator.get_kfac_blocks(examples)
            for layer_id, layer in \
                    self.generator.layer_collection.layers.items():
                a, g = kfac_blocks[layer_id]
                evals_a, evecs_a = torch.linalg.eigh(a)
                evals_g, evecs_g = torch.linalg.eigh(g)
                evecs[layer_id] = (evecs_a, evecs_g)
                diags[layer_id] = kronecker(evals_g.view(-1, 1),
                                            evals_a.view(-1, 1))
                del a, g, kfac_blocks[layer_id]
            self.data = (evecs, diags)
        else:
            self.data = data
//  [python]
def forward(self, A):
        w, v = torch.symeig(A, eigenvectors=True)

        self.save_for_backward(w, v)
        return w, v
//  [python]
def forward(self, A):
        w, v = torch.linalg.eigh(A)

        self.save_for_backward(w, v)
        return w, v
//  [python]
def forward(ctx, input):
        lambda_, v = torch.symeig(input, eigenvectors=True, upper=True)
        ctx.save_for_backward(input, lambda_, v)
        return lambda_, v
//  [python]
def forward(ctx, input):
        lambda_, v = torch.linalg.eigh(input, UPLO='U')
        ctx.save_for_backward(input, lambda_, v)
        return lambda_, v
//  [python]
def forward(ctx, A):
    eival, eivec = torch.symeig(A, eigenvectors=True)
    ctx.save_for_backward(eival, eivec)
    return eival, eivec
//  [python]
def forward(ctx, A):
    eival, eivec = torch.linalg.eigh(A)
    ctx.save_for_backward(eival, eivec)
    return eival, eivec
//  [python]
def A_vec_to_quat(A_vec):
    A = convert_Avec_to_A(A_vec)
    if A.dim() < 3:
        A = A.unsqueeze(dim=0)
    _, evs = torch.symeig(A, eigenvectors=True)
    return evs[:,:,0].squeeze()
//  [python]
def A_vec_to_quat(A_vec):
    A = convert_Avec_to_A(A_vec)
    if A.dim() < 3:
        A = A.unsqueeze(dim=0)
    _, evs = torch.linalg.eigh(A)
    return evs[:,:,0].squeeze()
//  [python]
def __init__(self):
		self.n_vertices = np.array([51, 51])
		self.n_factors = len(self.n_vertices)
		self.suggested_init = torch.Tensor(self.n_vertices).long().unsqueeze(0) / 2
		for _ in range(1, 2):
			random_init = torch.cat([torch.randint(0, int(elm), (1, 1)) for elm in self.n_vertices], dim=1)
			self.suggested_init = torch.cat([self.suggested_init, random_init], dim=0)
		self.adjacency_mat = []
		self.fourier_freq = []
		self.fourier_basis = []
		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)
			adjmat *= (n_v - 1.0)
			self.adjacency_mat.append(adjmat)
			degmat = torch.sum(adjmat, dim=0)
			laplacian = (torch.diag(degmat) - adjmat)
			eigval, eigvec = torch.symeig(laplacian, eigenvectors=True)
			self.fourier_freq.append(eigval)
			self.fourier_basis.append(eigvec)
//  [python]
def __init__(self):
		self.n_vertices = np.array([51, 51])
		self.n_factors = len(self.n_vertices)
		self.suggested_init = torch.Tensor(self.n_vertices).long().unsqueeze(0) / 2
		for _ in range(1, 2):
			random_init = torch.cat([torch.randint(0, int(elm), (1, 1)) for elm in self.n_vertices], dim=1)
			self.suggested_init = torch.cat([self.suggested_init, random_init], dim=0)
		self.adjacency_mat = []
		self.fourier_freq = []
		self.fourier_basis = []
		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)
			adjmat *= (n_v - 1.0)
			self.adjacency_mat.append(adjmat)
			degmat = torch.sum(adjmat, dim=0)
			laplacian = (torch.diag(degmat) - adjmat)
			eigval, eigvec = torch.linalg.eigh(laplacian)
			self.fourier_freq.append(eigval)
			self.fourier_basis.append(eigvec)
//  [python]
def global_PCA_alignment(stacked_points, batch_lengths):

    i0 = 0
    covariances = []
    centroids = []
    for b_i, length in enumerate(batch_lengths):

        # Extract point cloud from batch
        points = stacked_points[i0:i0 + length]

        # Center points on centroid [n_points, dim]
        centroid = torch.mean(points, dim=0, keepdim=True)
        points0 = points - centroid

        # Get covariance matrix [dim, n_points] x [n_points, dim]
        covariances.append(torch.matmul(points0.t(), points0))
        centroids.append(centroid)

        # Increment for next cloud
        i0 += length

    covariances = torch.stack(covariances, dim=0).cpu()
    centroids = torch.cat(centroids, dim=0)

    # Get eigenvalues and eigenvectors of covariances matrices [n_points, dim] and [n_points, dim, dim]
    eigen_values, eigen_vectors = torch.symeig(covariances, eigenvectors=True)
    eigen_vectors = eigen_vectors.to(stacked_points.device)

    # Correct rotations with centroids
    rotated_centroids = torch.matmul(centroids.unsqueeze(1), eigen_vectors)
    corrections = (rotated_centroids < 0).type(torch.float32) * 2 - 1

    # Get rotation matrices. if X is neighbors, then R.T x X is the aligned neighbors [n_points, dim, dim]
    rotations = eigen_vectors * corrections

    # Correct the orientation to avoid symmetries
    rotations[:, :, 2] = torch.cross(rotations[:, :, 0], rotations[:, :, 1])

    return rotations
//  [python]
def global_PCA_alignment(stacked_points, batch_lengths):

    i0 = 0
    covariances = []
    centroids = []
    for b_i, length in enumerate(batch_lengths):

        # Extract point cloud from batch
        points = stacked_points[i0:i0 + length]

        # Center points on centroid [n_points, dim]
        centroid = torch.mean(points, dim=0, keepdim=True)
        points0 = points - centroid

        # Get covariance matrix [dim, n_points] x [n_points, dim]
        covariances.append(torch.matmul(points0.t(), points0))
        centroids.append(centroid)

        # Increment for next cloud
        i0 += length

    covariances = torch.stack(covariances, dim=0).cpu()
    centroids = torch.cat(centroids, dim=0)

    # Get eigenvalues and eigenvectors of covariances matrices [n_points, dim] and [n_points, dim, dim]
    eigen_values, eigen_vectors = torch.linalg.eigh(covariances)
    eigen_vectors = eigen_vectors.to(stacked_points.device)

    # Correct rotations with centroids
    rotated_centroids = torch.matmul(centroids.unsqueeze(1), eigen_vectors)
    corrections = (rotated_centroids < 0).type(torch.float32) * 2 - 1

    # Get rotation matrices. if X is neighbors, then R.T x X is the aligned neighbors [n_points, dim, dim]
    rotations = eigen_vectors * corrections

    # Correct the orientation to avoid symmetries
    rotations[:, :, 2] = torch.cross(rotations[:, :, 0], rotations[:, :, 1])

    return rotations
//  [python]
def __init__(self):
		self.n_vertices = np.array([51] * 6)
		self.n_factors = len(self.n_vertices)
		self.suggested_init = torch.Tensor(self.n_vertices).long().unsqueeze(0) / 2
		for _ in range(1, 2):
			random_init = torch.cat([torch.randint(0, int(elm), (1, 1)) for elm in self.n_vertices], dim=1)
			self.suggested_init = torch.cat([self.suggested_init, random_init], dim=0)
		self.adjacency_mat = []
		self.fourier_freq = []
		self.fourier_basis = []
		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)
			adjmat *= (n_v - 1.0)
			self.adjacency_mat.append(adjmat)
			wgtsum = torch.sum(adjmat, dim=0)
			laplacian = (torch.diag(wgtsum) - adjmat)
			eigval, eigvec = torch.symeig(laplacian, eigenvectors=True)
			self.fourier_freq.append(eigval)
			self.fourier_basis.append(eigvec)
//  [python]
def __init__(self):
		self.n_vertices = np.array([51] * 6)
		self.n_factors = len(self.n_vertices)
		self.suggested_init = torch.Tensor(self.n_vertices).long().unsqueeze(0) / 2
		for _ in range(1, 2):
			random_init = torch.cat([torch.randint(0, int(elm), (1, 1)) for elm in self.n_vertices], dim=1)
			self.suggested_init = torch.cat([self.suggested_init, random_init], dim=0)
		self.adjacency_mat = []
		self.fourier_freq = []
		self.fourier_basis = []
		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)
			adjmat *= (n_v - 1.0)
			self.adjacency_mat.append(adjmat)
			wgtsum = torch.sum(adjmat, dim=0)
			laplacian = (torch.diag(wgtsum) - adjmat)
			eigval, eigvec = torch.linalg.eigh(laplacian)
			self.fourier_freq.append(eigval)
			self.fourier_basis.append(eigvec)
//  [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 _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, 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.linalg.eigh(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 solve_wahba_fast(A, compute_gap=False):
    """
    Use a fast eigenvalue solution to the dual of the 'generalized Wahba' problem to solve the primal.
    :param A: quadratic cost matrix
    :param redundant_constraints: boolean indicating whether to use redundand constraints
    :return: Optimal q, optimal dual var. nu, time to solve, duality gap
    """
    #start = time.time()
    # Returns (b,n) and (b,n,n) tensors
    nus, qs = torch.symeig(A, eigenvectors=True)
    nu_min, nu_argmin = torch.min(nus, 1)# , keepdim=False, out=None)
    q_opt = qs[torch.arange(A.shape[0]), :, nu_argmin]
    q_opt = q_opt*(torch.sign(q_opt[:, 3]).unsqueeze(1))
    nu_opt = -1.*nu_min.unsqueeze(1)
    if compute_gap:
        p = torch.einsum('bn,bnm,bm->b', q_opt, A, q_opt).unsqueeze(1)
        gap = p + nu_opt
        return q_opt, nu_opt, gap
    return q_opt, nu_opt
//  [python]
def solve_wahba_fast(A, compute_gap=False):
    """
    Use a fast eigenvalue solution to the dual of the 'generalized Wahba' problem to solve the primal.
    :param A: quadratic cost matrix
    :param redundant_constraints: boolean indicating whether to use redundand constraints
    :return: Optimal q, optimal dual var. nu, time to solve, duality gap
    """
    #start = time.time()
    # Returns (b,n) and (b,n,n) tensors
    nus, qs = torch.linalg.eigh(A)
    nu_min, nu_argmin = torch.min(nus, 1)# , keepdim=False, out=None)
    q_opt = qs[torch.arange(A.shape[0]), :, nu_argmin]
    q_opt = q_opt*(torch.sign(q_opt[:, 3]).unsqueeze(1))
    nu_opt = -1.*nu_min.unsqueeze(1)
    if compute_gap:
        p = torch.einsum('bn,bnm,bm->b', q_opt, A, q_opt).unsqueeze(1)
        gap = p + nu_opt
        return q_opt, nu_opt, gap
    return q_opt, nu_opt
//  [python]
def compute_rotation_matrix_from_10d(x):
    A = convert_Avec_to_A(x)
    d = A.device
    _, evs = torch.symeig(A.cpu(), eigenvectors=True)
    evs = evs.to(d)
    q = evs[:,:,0]
    return compute_rotation_matrix_from_quaternion(q,n_flag=False)
//  [python]
def compute_rotation_matrix_from_10d(x):
    A = convert_Avec_to_A(x)
    d = A.device
    _, evs = torch.linalg.eigh(A.cpu())
    evs = evs.to(d)
    q = evs[:,:,0]
    return compute_rotation_matrix_from_quaternion(q,n_flag=False)
//  [python]
def add_eigen(data):
    with torch.no_grad():
        point_mat = data.unsqueeze(2) - data.unsqueeze(1) 

        A = torch.exp(-torch.sum(point_mat.square(), dim=3))
        D = torch.diag_embed(torch.sum(A, dim=2))
        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 torch.cat((data, v), dim=2)
//  [python]
def add_eigen(data):
    with torch.no_grad():
        point_mat = data.unsqueeze(2) - data.unsqueeze(1) 

        A = torch.exp(-torch.sum(point_mat.square(), dim=3))
        D = torch.diag_embed(torch.sum(A, dim=2))
        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 torch.cat((data, v), dim=2)
//  [python]
def noise_to_sty_code(self, noise):
		sty_cov = self.sty.get_normalized_cov().cpu()
		sty_eigval, sty_eigvec = torch.symeig(sty_cov, eigenvectors = True)
		coef = torch.matmul(noise, sty_eigvec.to(noise))
		return torch.matmul(coef.mul(sty_eigval.sqrt().unsqueeze(0).to(noise)), sty_eigvec.transpose(0, 1).to(noise))
//  [python]
def noise_to_sty_code(self, noise):
		sty_eigval, sty_eigvec = torch.linalg.eigh(self.sty.get_normalized_cov().cpu())
		coef = torch.matmul(noise, sty_eigvec.to(noise))
		return torch.matmul(coef.mul(sty_eigval.sqrt().unsqueeze(0).to(noise)), sty_eigvec.transpose(0, 1).to(noise))
//  [python]
def compute_patch_whitening_statistics(train_set):

    def cov(X):
        X = X/np.sqrt(X.size(0) - 1)
        return X.t() @ X

    def patches(data, patch_size=(3, 3), dtype=torch.float32):
        h, w = patch_size
        c = data.size(1)
        return data.unfold(2, h, 1).unfold(3, w, 1).transpose(1, 3).reshape(-1, c, h, w).to(dtype)

    def eigens(patches):
        n, c, h, w = patches.shape
        covariance = cov(patches.reshape(n, c*h*w))
        eigen_values, eigen_vectors = torch.symeig(covariance, eigenvectors=True)
        return eigen_values.flip(0), eigen_vectors.t().reshape(c*h*w, c, h, w).flip(0)

    eigen_values, eigen_vectors = eigens(
        patches(train_set['data'][:10000, :, 4:-4, 4:-4])
    ) # center crop to remove padding
    return eigen_values, eigen_vectors
//  [python]
def compute_patch_whitening_statistics(train_set):

    def cov(X):
        X = X/np.sqrt(X.size(0) - 1)
        return X.t() @ X

    def patches(data, patch_size=(3, 3), dtype=torch.float32):
        h, w = patch_size
        c = data.size(1)
        return data.unfold(2, h, 1).unfold(3, w, 1).transpose(1, 3).reshape(-1, c, h, w).to(dtype)

    def eigens(patches):
        n, c, h, w = patches.shape
        covariance = cov(patches.reshape(n, c*h*w))
        eigen_values, eigen_vectors = torch.linalg.eigh(covariance)
        return eigen_values.flip(0), eigen_vectors.t().reshape(c*h*w, c, h, w).flip(0)

    eigen_values, eigen_vectors = eigens(
        patches(train_set['data'][:10000, :, 4:-4, 4:-4])
    ) # center crop to remove padding
    return eigen_values, eigen_vectors
//  [python]
def symsqrt(self, a, cond=None, return_rank=False):
		"""Computes the symmetric square root of a positive definite matrix
		Thanks to 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[a.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(self, a, cond=None, return_rank=False):
		"""Computes the symmetric square root of a positive definite matrix
		Thanks to 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[a.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 _generalized_eigenvalue_decomposition(a, b):
    cholesky = stable_cholesky(b)
    inv_cholesky = torch.inverse(cholesky)
    # Compute C matrix L⁻1 A L^-T
    cmat = inv_cholesky @ a @ inv_cholesky.conj().transpose(-1, -2)
    # Performing the eigenvalue decomposition
    e_val, e_vec = torch.symeig(cmat, eigenvectors=True)
    # Collecting the eigenvectors
    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)
    # Compute C matrix L⁻1 A L^-T
    cmat = inv_cholesky @ a @ inv_cholesky.conj().transpose(-1, -2)
    # Performing the eigenvalue decomposition
    e_val, e_vec = torch.linalg.eigh(cmat)
    # Collecting the eigenvectors
    e_vec = torch.matmul(inv_cholesky.conj().transpose(-1, -2), e_vec)
    return e_val, e_vec
//  [python]
def max_eigenvalue_constraint_torch(x, maximum_eigenvalue):
    """
    This function defines an inequality constraint on the maximum eigenvalue of a SPD matrix.
    The value returned by the function is positive if the inequality constraints is satisfied.

    Parameters
    ----------
    :param x: SPD matrix
    :param maximum_eigenvalue: maximum eigenvalue to satisfy the constraint

    Returns
    -------
    :return: difference between maximum tolerated eigenvalue and maximum eigenvalue of x
    """
    eigenvalues = torch.symeig(x, eigenvectors=True).eigenvalues  # Eigenvalue True necessary for derivation
    return maximum_eigenvalue - eigenvalues.max()
//  [python]
def max_eigenvalue_constraint_torch(x, maximum_eigenvalue):
    eigenvalues = torch.linalg.eigh(x).eigenvalues
    return maximum_eigenvalue - eigenvalues.max()
//  [python]
def forward(
        self,
        mix: torch.Tensor,
        target_scm: torch.Tensor,
        noise_scm: torch.Tensor,
    ):
        r"""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)
        """
        # TODO: Implement several RTF estimation strategies, and choose one here, or expose all.
        # Get relative transfer function (1st PCA of Σss)
        e_val, e_vec = torch.symeig(target_scm.permute(0, 3, 1, 2), eigenvectors=True)
        rtf_vect = e_vec[..., -1]  # bfm
        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,
    ):
        r"""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)
        """
        # TODO: Implement several RTF estimation strategies, and choose one here, or expose all.
        # Get relative transfer function (1st PCA of Σss)
        e_val, e_vec = torch.linalg.eigh(target_scm.permute(0, 3, 1, 2))
        rtf_vect = e_vec[..., -1]  # bfm
        return self.from_rtf_vect(mix=mix, rtf_vec=rtf_vect.transpose(-1, -2), noise_scm=noise_scm)
//  [python]
def min_eigenvalue_constraint_torch(x, minimum_eigenvalue):
    """
    This function defines an inequality constraint on the minimum eigenvalue of a SPD matrix.
    The value returned by the function is positive if the inequality constraints is satisfied.

    Parameters
    ----------
    :param x: SPD matrix
    :param minimum_eigenvalue: minimum eigenvalue to satisfy the constraint

    Returns
    -------
    :return: difference between minimum eigenvalue of x and minimum tolerated eigenvalue
    """
    eigenvalues = torch.symeig(x, eigenvectors=True).eigenvalues  # Eigenvalue True necessary for derivation
    return eigenvalues.min() - minimum_eigenvalue
//  [python]
def min_eigenvalue_constraint_torch(x, minimum_eigenvalue):
    eigenvalues = torch.linalg.eigh(x).eigenvalues
    return eigenvalues.min() - minimum_eigenvalue
//  [python]
def covsqrt_mean(feature, inverse=False, tolerance=1e-14):
    # I referenced the default svd tolerance value in matlab.

    b, c, h, w = feature.size()

    mean = torch.mean(feature.view(b, c, -1), dim=2, keepdim=True)
    zeromean = feature.view(b, c, -1) - mean
    cov = torch.bmm(zeromean, zeromean.transpose(1, 2))

    evals, evects = torch.symeig(cov, eigenvectors=True)
    
    p = 0.5
    if inverse:
        p *= -1

    covsqrt = []
    for i in range(b):
        k = 0
        for j in range(c):
            if evals[i][j] > tolerance:
                k = j
                break
        covsqrt.append(torch.mm(evects[i][:, k:],
                            torch.mm(evals[i][k:].pow(p).diag_embed(),
                                     evects[i][:, k:].t())).unsqueeze(0))
    covsqrt = torch.cat(covsqrt, dim=0)

    return covsqrt, mean
//  [python]
def covsqrt_mean(feature, inverse=False, tolerance=1e-14):
    # I referenced the default svd tolerance value in matlab.

    b, c, h, w = feature.size()

    mean = torch.mean(feature.view(b, c, -1), dim=2, keepdim=True)
    zeromean = feature.view(b, c, -1) - mean
    cov = torch.bmm(zeromean, zeromean.transpose(1, 2))

    evals, evects = torch.linalg.eigh(cov)
    
    p = 0.5
    if inverse:
        p *= -1

    covsqrt = []
    for i in range(b):
        k = 0
        for j in range(c):
            if evals[i][j] > tolerance:
                k = j
                break
        covsqrt.append(torch.mm(evects[i][:, k:],
                            torch.mm(evals[i][k:].pow(p).diag_embed(),
                                     evects[i][:, k:].t())).unsqueeze(0))
    covsqrt = torch.cat(covsqrt, dim=0)

    return covsqrt, mean
//  [python]
def finalize(self):
        """
        Finalize training with LU factorization or Pseudo-inverse
        """
        # Reshape average
        xTx, avg, tlen = self._fix(self.xTx, self.xTx_avg, self.tlen)

        # Reshape
        self.avg = avg.unsqueeze(0)

        # We need more observations than variables
        if self.tlen < self.input_dim:
            raise Exception(u"The number of observations ({}) is larger than  the number of input variables ({})".format(self.tlen, self.input_dim))
        # end if

        # Total variance
        total_var = torch.diag(xTx).sum()

        # Compute and sort eigenvalues
        d, v = torch.symeig(xTx, eigenvectors=True)

        # Check for negative eigenvalues
        if float(d.min()) < 0:
            # raise Exception(u"Got negative eigenvalues ({}). You may either set output_dim to be smaller".format(d))
            pass
        # end if

        # Indexes
        indexes = range(d.size(0)-1, -1, -1)

        # Sort by descending order
        d = torch.take(d, Variable(torch.LongTensor(indexes)))
        v = v[:, indexes]

        # Explained covariance
        self.explained_variance = torch.sum(d) / total_var

        # Store eigenvalues
        self.d = d[:self.output_dim]

        # Store eigenvectors
        self.v = v[:, :self.output_dim]

        # Total variance
        self.total_variance = total_var

        # Stop training
        self.train(False)
//  [python]
def finalize(self):
        """
        Finalize training with LU factorization or Pseudo-inverse
        """
        # Reshape average
        xTx, avg, tlen = self._fix(self.xTx, self.xTx_avg, self.tlen)

        # Reshape
        self.avg = avg.unsqueeze(0)

        # We need more observations than variables
        if self.tlen < self.input_dim:
            raise Exception(u"The number of observations ({}) is larger than  the number of input variables ({})".format(self.tlen, self.input_dim))
        # end if

        # Total variance
        total_var = torch.diag(xTx).sum()

        # Compute and sort eigenvalues
        d, v = torch.linalg.eigh(xTx)

        # Check for negative eigenvalues
        if float(d.min()) < 0:
            # raise Exception(u"Got negative eigenvalues ({}). You may either set output_dim to be smaller".format(d))
            pass
        # end if

        # Indexes
        indexes = range(d.size(0)-1, -1, -1)

        # Sort by descending order
        d = torch.take(d, Variable(torch.LongTensor(indexes)))
        v = v[:, indexes]

        # Explained covariance
        self.explained_variance = torch.sum(d) / total_var

        # Store eigenvalues
        self.d = d[:self.output_dim]

        # Store eigenvectors
        self.v = v[:, :self.output_dim]

        # Total variance
        self.total_variance = total_var

        # Stop training
        self.train(False)
//  [python]
def covsqrt_mean(feature, inverse=False, tolerance=1e-14):
    b, c, h, w = feature.size()
    mean = torch.mean(feature.view(b, c, -1), dim=2, keepdim=True)
    zeromean = feature.view(b, c, -1) - mean
    cov = torch.bmm(zeromean, zeromean.transpose(1, 2))
    evals, evects = torch.symeig(cov, eigenvectors=True)
    p = 0.5
    if inverse:
        p *= -1
    covsqrt = []
    for i in range(b):
        k = 0
        for j in range(c):
            if evals[i][j] > tolerance:
                k = j
                break
        covsqrt.append(torch.mm(evects[i][:, k:], torch.mm(evals[i][k:].pow(p).diag_embed(), evects[i][:, k:].t())).unsqueeze(0))
    covsqrt = torch.cat(covsqrt, dim=0)
    return covsqrt, mean
//  [python]
def covsqrt_mean(feature, inverse=False, tolerance=1e-14):
    b, c, h, w = feature.size()
    mean = torch.mean(feature.view(b, c, -1), dim=2, keepdim=True)
    zeromean = feature.view(b, c, -1) - mean
    cov = torch.bmm(zeromean, zeromean.transpose(1, 2))
    evals, evects = torch.linalg.eigh(cov)
    p = 0.5
    if inverse:
        p *= -1
    covsqrt = []
    for i in range(b):
        k = 0
        for j in range(c):
            if evals[i][j] > tolerance:
                k = j
                break
        covsqrt.append(torch.mm(evects[i][:, k:], torch.mm(evals[i][k:].pow(p).diag_embed(), evects[i][:, k:].t())).unsqueeze(0))
    covsqrt = torch.cat(covsqrt, dim=0)
    return covsqrt, mean
//  [python]
def post_processing_init_spd_torch(x_vec, min_eigenvalue, max_eigenvalue):
    """
    This function post-processes a symmetric matrix, so that its eigenvalues lie in the defined bounds.

    Parameters
    ----------
    :param x_vec: symmetric matrix in Mandel vector form
    :param min_eigenvalue: minimum eigenvalue to satisfy the constraint
    :param max_eigenvalue: maximum eigenvalue to satisfy the constraint

    Returns
    -------
    :return: symmetric matrix in Mandel vector form
    """
    x = vector_to_symmetric_matrix_mandel_torch(x_vec)
    init_shape = x.shape
    x = x.view(-1, x.shape[-2], x.shape[-1])

    for n in range(x.shape[0]):
        eigdec = torch.symeig(x[n], eigenvectors=True)
        eigvals = eigdec.eigenvalues
        eigvecs = eigdec.eigenvectors

        eigvals[eigvals <= min_eigenvalue] = min_eigenvalue  # Minimum eigenvalue constraint
        eigvals[eigvals > max_eigenvalue] = max_eigenvalue  # Max eigenvalue constraint

        x[n] = torch.mm(torch.mm(eigvecs, torch.diag(eigvals)), torch.inverse(eigvecs))

    x = x.view(init_shape)
    return symmetric_matrix_to_vector_mandel_torch(x)
//  [python]
def post_processing_init_spd_torch(x_vec, min_eigenvalue, max_eigenvalue):
    """
    This function post-processes a symmetric matrix, so that its eigenvalues lie in the defined bounds.

    Parameters
    ----------
    :param x_vec: symmetric matrix in Mandel vector form
    :param min_eigenvalue: minimum eigenvalue to satisfy the constraint
    :param max_eigenvalue: maximum eigenvalue to satisfy the constraint

    Returns
    -------
    :return: symmetric matrix in Mandel vector form
    """
    x = vector_to_symmetric_matrix_mandel_torch(x_vec)
    init_shape = x.shape
    x = x.view(-1, x.shape[-2], x.shape[-1])

    for n in range(x.shape[0]):
        eigdec = torch.linalg.eigh(x[n])
        eigvals = eigdec.eigenvalues
        eigvecs = eigdec.eigenvectors

        eigvals[eigvals <= min_eigenvalue] = min_eigenvalue  # Minimum eigenvalue constraint
        eigvals[eigvals > max_eigenvalue] = max_eigenvalue  # Max eigenvalue constraint

        x[n] = torch.mm(torch.mm(eigvecs, torch.diag(eigvals)), torch.inverse(eigvecs))

    x = x.view(init_shape)
    return symmetric_matrix_to_vector_mandel_torch(x)