jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def negative_log_likelihood(self, space_coordinates, time_coordinates, stData):
        """
        Calculates the negative log-likelihood of the MultiTaskGP_IMC model.

        Args:
            space_coordinates (torch.Tensor): Tensor containing the space coordinates.
            time_coordinates (torch.Tensor): Tensor containing the time coordinates.
            stData (torch.Tensor): Tensor containing the data.

        Returns:
            The negative log-likelihood value.
        """
        latlong_K = self.latlong_kernel(space_coordinates[:, 0:2], space_coordinates[:, 0:2])
        elevation_K = self.elevation_kernel(space_coordinates[:, 2:3], space_coordinates[:, 2:3])

        spatial_K = latlong_K * elevation_K + torch.eye(latlong_K.size(0)) * JITTER
        temporal_K = self.temporal_kernel(time_coordinates, time_coordinates) + torch.eye(time_coordinates.size(0)) * JITTER

        # Matrix Eigendecomposition
        eigen_value_s, eigen_vector_s = torch.symeig(spatial_K, eigenvectors=True)
        eigen_value_t, eigen_vector_t = torch.symeig(temporal_K, eigenvectors=True)

        eigen_vector_st = kronecker(eigen_vector_t, eigen_vector_s)
        eigen_value_st = kronecker(eigen_value_t.view(-1, 1), eigen_value_s.view(-1, 1)).view(-1)
        eigen_value_st_plus_noise_inverse = 1. / (eigen_value_st + torch.exp(self.log_noise_variance))

        Lambda_st = eigen_value_st_plus_noise_inverse.diag_embed()
        A = torch.flatten(eigen_vector_t.transpose(-2, -1) @ stData.transpose(-2, -1) @ eigen_vector_s).unsqueeze(-1)        
        self.alpha = Lambda_st @ A
        self.Lambda_st = Lambda_st
        self.eigen_vector_st = eigen_vector_st

        nll = 0
        nll += 0.5 * (eigen_value_st + torch.exp(self.log_noise_variance)).log().sum()
        nll += 0.5 * (A.transpose(-2, -1) @ self.alpha).sum()
        return nll
//  [python]
def negative_log_likelihood(self, space_coordinates, time_coordinates, stData):
        """
        Calculates the negative log-likelihood of the MultiTaskGP_IMC model.

        Args:
            space_coordinates (torch.Tensor): Tensor containing the space coordinates.
            time_coordinates (torch.Tensor): Tensor containing the time coordinates.
            stData (torch.Tensor): Tensor containing the data.

        Returns:
            The negative log-likelihood value.
        """
        latlong_K = self.latlong_kernel(space_coordinates[:, 0:2], space_coordinates[:, 0:2])
        elevation_K = self.elevation_kernel(space_coordinates[:, 2:3], space_coordinates[:, 2:3])

        spatial_K = latlong_K * elevation_K + torch.eye(latlong_K.size(0)) * JITTER
        temporal_K = self.temporal_kernel(time_coordinates, time_coordinates) + torch.eye(time_coordinates.size(0)) * JITTER

        # Matrix Eigendecomposition
        eigen_value_s, eigen_vector_s = torch.linalg.eigh(spatial_K, UPLO='L')
        eigen_value_t, eigen_vector_t = torch.linalg.eigh(temporal_K, UPLO='L')

        eigen_vector_st = kronecker(eigen_vector_t, eigen_vector_s)
        eigen_value_st = kronecker(eigen_value_t.view(-1, 1), eigen_value_s.view(-1, 1)).view(-1)
        eigen_value_st_plus_noise_inverse = 1. / (eigen_value_st + torch.exp(self.log_noise_variance))

        Lambda_st = eigen_value_st_plus_noise_inverse.diag_embed()
        A = torch.flatten(eigen_vector_t.transpose(-2, -1) @ stData.transpose(-2, -1) @ eigen_vector_s).unsqueeze(-1)        
        self.alpha = Lambda_st @ A
        self.Lambda_st = Lambda_st
        self.eigen_vector_st = eigen_vector_st

        nll = 0
        nll += 0.5 * (eigen_value_st + torch.exp(self.log_noise_variance)).log().sum()
        nll += 0.5 * (A.transpose(-2, -1) @ self.alpha).sum()
        return nll
//  [python]
def test_general_laplacian_builder_with_deg_normalisation_produces_valid_laplacian(graph_id, d):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = GeneralLaplacianBuilder(size, data.edge_index, d, deg_normalised=True)
    model.eval()

    maps = torch.FloatTensor(size=(data.edge_index.size(1), d, d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*d, size*d), dtype=torch.float64).to_dense()
    Q, _ = torch.symeig(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.allclose(L, L.T, atol=eps)
    assert torch.all(Q >= 0 - eps), Q
    assert torch.all(Q <= 2.)
//  [python]
def test_general_laplacian_builder_with_deg_normalisation_produces_valid_laplacian(graph_id, d):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = GeneralLaplacianBuilder(size, data.edge_index, d, deg_normalised=True)
    model.eval()

    maps = torch.FloatTensor(size=(data.edge_index.size(1), d, d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*d, size*d), dtype=torch.float64).to_dense()
    Q, _ = torch.linalg.eigh(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.allclose(L, L.T, atol=eps)
    assert torch.all(Q >= 0 - eps), Q
    assert torch.all(Q <= 2.)
//  [python]
def test_general_laplacian_builder_produces_valid_laplacian(graph_id, d, normalised, augmented):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = GeneralLaplacianBuilder(size, data.edge_index, d, normalised=normalised)
    model.eval()
    maps = torch.FloatTensor(size=(data.edge_index.size(1), d, d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*d, size*d), dtype=torch.float64).to_dense()
    Q, _ = torch.symeig(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.allclose(L, L.T, atol=eps)
    assert torch.all(Q >= 0 - eps), Q
    if normalised:
        assert torch.all(Q <= 2.)

    L_dense = build_dense_laplacian(size, data.edge_index, maps.to(dtype=torch.float64), d, normalised=normalised)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_general_laplacian_builder_produces_valid_laplacian(graph_id, d, normalised, augmented):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = GeneralLaplacianBuilder(size, data.edge_index, d, normalised=normalised)
    model.eval()
    maps = torch.FloatTensor(size=(data.edge_index.size(1), d, d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*d, size*d), dtype=torch.float64).to_dense()
    Q, _ = torch.linalg.eigh(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.allclose(L, L.T, atol=eps)
    assert torch.all(Q >= 0 - eps), Q
    if normalised:
        assert torch.all(Q <= 2.)

    L_dense = build_dense_laplacian(size, data.edge_index, maps.to(dtype=torch.float64), d, normalised=normalised)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_general_laplacian_builder_produces_valid_laplacian_on_texas(d, normalised, augmented):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    data = get_dataset('texas')[0]
    size = data.x.size(0)

    # Construct the model
    model = GeneralLaplacianBuilder(size, data.edge_index, d, normalised=normalised)
    model.eval()

    maps = torch.FloatTensor(size=(data.edge_index.size(1), d, d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*d, size*d), dtype=torch.float64).to_dense()
    Q, _ = torch.symeig(L)

    # Check that L is semi-positive definite.
    eps = 1e-5
    assert torch.allclose(L, L.T, atol=eps)
    assert torch.all(Q >= 0 - eps), Q
    if normalised:
        assert torch.all(Q <= 2.)

    L_dense = build_dense_laplacian(size, data.edge_index, maps, d, normalised=normalised)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_general_laplacian_builder_produces_valid_laplacian_on_texas(d, normalised, augmented):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    data = get_dataset('texas')[0]
    size = data.x.size(0)

    # Construct the model
    model = GeneralLaplacianBuilder(size, data.edge_index, d, normalised=normalised)
    model.eval()

    maps = torch.FloatTensor(size=(data.edge_index.size(1), d, d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*d, size*d), dtype=torch.float64).to_dense()
    Q, _ = torch.linalg.eigh(L)

    # Check that L is semi-positive definite.
    eps = 1e-5
    assert torch.allclose(L, L.T, atol=eps)
    assert torch.all(Q >= 0 - eps), Q
    if normalised:
        assert torch.all(Q <= 2.)

    L_dense = build_dense_laplacian(size, data.edge_index, maps, d, normalised=normalised)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_general_laplacian_builder_with_fixed_maps_produces_valid_laplacian(graph_id, d, normalised, augmented):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = GeneralLaplacianBuilder(size, data.edge_index, d, normalised=normalised, add_hp=True, add_lp=True)
    model.eval()

    maps = torch.FloatTensor(size=(data.edge_index.size(1), d, d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*(d+2), size*(d+2)), dtype=torch.float64).to_dense()
    Q, _ = torch.symeig(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.allclose(L, L.T, atol=eps)
    assert torch.all(Q >= 0 - eps), Q
    if normalised:
        assert torch.all(Q <= 2.)

    values = [1.0, -1.0]
    L_dense = build_dense_laplacian(size, data.edge_index, maps.to(dtype=torch.float64),
                                    d=d, normalised=normalised, values=values)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_general_laplacian_builder_with_fixed_maps_produces_valid_laplacian(graph_id, d, normalised, augmented):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = GeneralLaplacianBuilder(size, data.edge_index, d, normalised=normalised, add_hp=True, add_lp=True)
    model.eval()

    maps = torch.FloatTensor(size=(data.edge_index.size(1), d, d)).uniform_(-1.0, 1.0)
    L, _ = model(maps)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*(d+2), size*(d+2)), dtype=torch.float64).to_dense()
    Q, _ = torch.linalg.eigh(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.allclose(L, L.T, atol=eps)
    assert torch.all(Q >= 0 - eps), Q
    if normalised:
        assert torch.all(Q <= 2.)

    values = [1.0, -1.0]
    L_dense = build_dense_laplacian(size, data.edge_index, maps.to(dtype=torch.float64),
                                    d=d, normalised=normalised, values=values)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def tensor_exp(self, t):#4dim
        # condition: t is symmetric!
        s, u = torch.symeig(t)
        return u @ torch.diag_embed(torch.exp(s)) @ u.permute(0, 1, 3, 2)
//  [python]
def tensor_exp(self, t):#4dim
        # condition: t is symmetric!
        s, u = torch.linalg.eigh(t)
        return u @ torch.diag_embed(torch.exp(s)) @ u.permute(0, 1, 3, 2)
//  [python]
def test_norm_connection_laplacian_produces_valid_laplacian(graph_id, d, orth_trans):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = NormConnectionLaplacianBuilder(size, edge_index=data.edge_index, d=d, orth_map=orth_trans)
    # Generate maps
    samples = data.edge_index.size(1)
    if orth_trans in ["householder", "euler"]:
        map_params = torch.FloatTensor(size=(samples, d * (d - 1) // 2)).uniform_(-1.0, 1.0)
    else:
        map_params = torch.FloatTensor(size=(samples, d * (d + 1) // 2)).uniform_(-1.0, 1.0)
    maps = model.orth_transform(map_params)

    L, _ = model(map_params)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*d, size*d), dtype=torch.float64).to_dense()
    Q, _ = torch.symeig(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.equal(L, L.T)
    assert torch.all(Q >= 0 - eps), Q
    assert torch.all(Q <= 2. + eps), Q

    L_dense = build_dense_laplacian(size, data.edge_index, maps, d, normalised=True)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_norm_connection_laplacian_produces_valid_laplacian(graph_id, d, orth_trans):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = NormConnectionLaplacianBuilder(size, edge_index=data.edge_index, d=d, orth_map=orth_trans)
    # Generate maps
    samples = data.edge_index.size(1)
    if orth_trans in ["householder", "euler"]:
        map_params = torch.FloatTensor(size=(samples, d * (d - 1) // 2)).uniform_(-1.0, 1.0)
    else:
        map_params = torch.FloatTensor(size=(samples, d * (d + 1) // 2)).uniform_(-1.0, 1.0)
    maps = model.orth_transform(map_params)

    L, _ = model(map_params)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*d, size*d), dtype=torch.float64).to_dense()
    Q, _ = torch.linalg.eigh(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.equal(L, L.T)
    assert torch.all(Q >= 0 - eps), Q
    assert torch.all(Q <= 2. + eps), Q

    L_dense = build_dense_laplacian(size, data.edge_index, maps, d, normalised=True)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def hosvd(x, dc=None, return_lms=False):
    if type(dc) is int:
        dc = [dc] * x.ndimension()
    u, lms = list(), list()
    for n in range(x.ndimension()):
        mat = reduced_matrix(x, n)
        lm_, u_ = t.linalg.eigh(mat)
        if (dc is not None) and (dc[n] < x.shape[n]):
            u_ = u_[:, -dc[n]:]
        u.append(u_)
        lms.append(lm_)
    core = tucker_product(x, u, conj=True, dim=0)
    if return_lms:
        return core, u, lms
    else:
        return core, u
//  [python]
def hosvd(x, dc=None, return_lms=False):
    if type(dc) is int:
        dc = [dc] * x.ndimension()
    u, lms = list(), list()
    for n in range(x.ndimension()):
        mat = reduced_matrix(x, n)
        lm_, u_ = tc.linalg.eigh(mat)
        if (dc is not None) and (dc[n] < x.shape[n]):
            u_ = u_[:, -dc[n]:]
        u.append(u_)
        lms.append(lm_)
    core = tucker_product(x, u, conj=True, dim=0)
    if return_lms:
        return core, u, lms
    else:
        return core, u
//  [python]
def test_norm_connection_laplacian_with_fixed_maps_produces_valid_laplacian(graph_id, d, orth_trans):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = NormConnectionLaplacianBuilder(
        size, edge_index=data.edge_index, d=d, orth_map=orth_trans, add_lp=True, add_hp=True)
    # Generate maps
    samples = data.edge_index.size(1)
    if orth_trans in ["householder", "euler"]:
        map_params = torch.FloatTensor(size=(samples, d * (d - 1) // 2)).uniform_(-1.0, 1.0)
    else:
        map_params = torch.FloatTensor(size=(samples, d * (d + 1) // 2)).uniform_(-1.0, 1.0)
    maps = model.orth_transform(map_params)

    L, _ = model(map_params)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*(d+2), size*(d+2)), dtype=torch.float64).to_dense()
    Q, _ = torch.symeig(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.equal(L, L.T)
    assert torch.all(Q >= 0 - eps), Q
    assert torch.all(Q <= 2. + eps), Q

    # Append fixed submaps to the restriction maps.
    values = [1.0, -1.0]
    L_dense = build_dense_laplacian(size, data.edge_index, maps, d, normalised=True, values=values)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_norm_connection_laplacian_with_fixed_maps_produces_valid_laplacian(graph_id, d, orth_trans):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = NormConnectionLaplacianBuilder(
        size, edge_index=data.edge_index, d=d, orth_map=orth_trans, add_lp=True, add_hp=True)
    # Generate maps
    samples = data.edge_index.size(1)
    if orth_trans in ["householder", "euler"]:
        map_params = torch.FloatTensor(size=(samples, d * (d - 1) // 2)).uniform_(-1.0, 1.0)
    else:
        map_params = torch.FloatTensor(size=(samples, d * (d + 1) // 2)).uniform_(-1.0, 1.0)
    maps = model.orth_transform(map_params)

    L, _ = model(map_params)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*(d+2), size*(d+2)), dtype=torch.float64).to_dense()
    Q, _ = torch.linalg.eigh(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.equal(L, L.T)
    assert torch.all(Q >= 0 - eps), Q
    assert torch.all(Q <= 2. + eps), Q

    # Append fixed submaps to the restriction maps.
    values = [1.0, -1.0]
    L_dense = build_dense_laplacian(size, data.edge_index, maps, d, normalised=True, values=values)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_norm_connection_laplacian_with_edge_weights_produces_valid_laplacian(graph_id, d, orth_trans):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = NormConnectionLaplacianBuilder(
        size, edge_index=data.edge_index, d=d, orth_map=orth_trans, add_lp=True, add_hp=True)
    # Generate maps
    samples = data.edge_index.size(1)
    if orth_trans in ["householder", "euler"]:
        map_params = torch.FloatTensor(size=(samples, d * (d - 1) // 2)).uniform_(-1.0, 1.0)
    else:
        map_params = torch.FloatTensor(size=(samples, d * (d + 1) // 2)).uniform_(-1.0, 1.0)

    edge_weights = lap.get_random_edge_weights(data.edge_index)
    # Use edge weights
    maps = model.orth_transform(map_params)

    L, _ = model(map_params, edge_weights)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*(d+2), size*(d+2)), dtype=torch.float64).to_dense()
    Q, _ = torch.symeig(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.equal(L, L.T)
    assert torch.all(Q >= 0 - eps), Q
    assert torch.all(Q <= 2. + eps), Q

    # Append fixed submaps to the restriction maps.
    values = [1.0, -1.0]
    L_dense = build_dense_laplacian(size, data.edge_index, maps, d, normalised=True, values=values,
                                    edge_weights=edge_weights)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def test_norm_connection_laplacian_with_edge_weights_produces_valid_laplacian(graph_id, d, orth_trans):
    # Fix the random seed
    torch.random.manual_seed(0)
    np.random.seed(0)

    # Build a test graph
    graph = nx.graph_atlas(graph_id)
    size = len(graph)
    data = from_networkx(graph)

    # Construct the model
    model = NormConnectionLaplacianBuilder(
        size, edge_index=data.edge_index, d=d, orth_map=orth_trans, add_lp=True, add_hp=True)
    # Generate maps
    samples = data.edge_index.size(1)
    if orth_trans in ["householder", "euler"]:
        map_params = torch.FloatTensor(size=(samples, d * (d - 1) // 2)).uniform_(-1.0, 1.0)
    else:
        map_params = torch.FloatTensor(size=(samples, d * (d + 1) // 2)).uniform_(-1.0, 1.0)

    edge_weights = lap.get_random_edge_weights(data.edge_index)
    # Use edge weights
    maps = model.orth_transform(map_params)

    L, _ = model(map_params, edge_weights)
    L = torch.sparse_coo_tensor(L[0], L[1], size=(size*(d+2), size*(d+2)), dtype=torch.float64).to_dense()
    Q, _ = torch.linalg.eigh(L)

    # Check that L is semi-positive definite.
    eps = 1e-6
    assert torch.equal(L, L.T)
    assert torch.all(Q >= 0 - eps), Q
    assert torch.all(Q <= 2. + eps), Q

    # Append fixed submaps to the restriction maps.
    values = [1.0, -1.0]
    L_dense = build_dense_laplacian(size, data.edge_index, maps, d, normalised=True, values=values,
                                    edge_weights=edge_weights)
    assert torch.allclose(L, L_dense, atol=eps)
//  [python]
def symmetrize_lattice(
    self: symmetries.Symmetries, Rbasis: torch.Tensor
) -> torch.Tensor:
    """Symmetrize lattice vectors `Rbasis` (3 x 3)."""
    # Compute symmetrized metric:
    metric = Rbasis.T @ Rbasis
    metric_sym = (self.rot.transpose(-2, -1) @ (metric @ self.rot)).mean(dim=0)
    # From transformation from matrix square-roots of metrics:
    E, V = torch.symeig(metric)
    E_sym, V_sym = torch.symeig(metric_sym)
    return (
        Rbasis
        @ (V @ ((1.0 / E.sqrt()).diag_embed() @ V.T))  # metric^(-1/2)
        @ (V_sym @ (E_sym.sqrt().diag_embed() @ V_sym.T))
    )
//  [python]
def symmetrize_lattice(
    self: symmetries.Symmetries, Rbasis: torch.Tensor
) -> torch.Tensor:
    """Symmetrize lattice vectors `Rbasis` (3 x 3)."""
    # Compute symmetrized metric:
    metric = Rbasis.T @ Rbasis
    metric_sym = (self.rot.transpose(-2, -1) @ (metric @ self.rot)).mean(dim=0)
    # From transformation from matrix square-roots of metrics:
    E, V = torch.linalg.eigh(metric)
    E_sym, V_sym = torch.linalg.eigh(metric_sym)
    return (
        Rbasis
        @ (V @ ((1.0 / E.sqrt()).diag_embed() @ V.T))  # metric^(-1/2)
        @ (V_sym @ (E_sym.sqrt().diag_embed() @ V_sym.T))
    )
//  [python]
def forward(ctx, a: Tensor, method: str = 'cond', factor: float = 1E-12
                ) -> Tuple[Tensor, Tensor]:
        """Calculate the eigenvalues and eigenvectors of a symmetric matrix.

        Finds the eigenvalues and eigenvectors of a real symmetric
        matrix using the torch.symeig function.

        Arguments:
            a: A real symmetric matrix whose eigenvalues & eigenvectors will
                be computed.
            method: Broadening method to used, available options are:

                    - "cond" for conditional broadening.
                    - "lorn" for Lorentzian broadening.

                See class doc-string for more info on these methods.
                [DEFAULT='cond']
            factor: Degree of broadening (broadening factor). [Default=1E-12]

        Returns:
            w: The eigenvalues, in ascending order.
            v: The eigenvectors.

        Notes:
            The ctx argument is auto-parsed by PyTorch & is used to pass data
            from the .forward() method to the .backward() method. This is not
            normally described in the docstring but has been done here to act
            as an example.

        Warnings:
            Under no circumstances should `factor` be a torch.tensor entity.
            The `method` and `factor` parameters MUST be passed as positional
            arguments and NOT keyword arguments.

        """
        # Check that the method is of a known type
        if method not in _SymEigB.KNOWN_METHODS:
            raise ValueError('Unknown broadening method selected.')

        # Compute eigen-values & vectors using torch.symeig.
        w, v = torch.symeig(a)

        # Save tensors that will be needed in the backward pass
        ctx.save_for_backward(w, v)

        # Save the broadening factor and the selected broadening method.
        ctx.bf, ctx.bm = factor, method

        # Store dtype/device to prevent dtype/device mixing
        ctx.dtype, ctx.device = a.dtype, a.device

        # Return the eigenvalues and eigenvectors
        return w, v
//  [python]
def forward(ctx, a: Tensor, method: str = 'cond', factor: float = 1E-12
                ) -> Tuple[Tensor, Tensor]:
        """Calculate the eigenvalues and eigenvectors of a symmetric matrix.

        Finds the eigenvalues and eigenvectors of a real symmetric
        matrix using the torch.symeig function.

        Arguments:
            a: A real symmetric matrix whose eigenvalues & eigenvectors will
                be computed.
            method: Broadening method to used, available options are:

                    - "cond" for conditional broadening.
                    - "lorn" for Lorentzian broadening.

                See class doc-string for more info on these methods.
                [DEFAULT='cond']
            factor: Degree of broadening (broadening factor). [Default=1E-12]

        Returns:
            w: The eigenvalues, in ascending order.
            v: The eigenvectors.

        Notes:
            The ctx argument is auto-parsed by PyTorch & is used to pass data
            from the .forward() method to the .backward() method. This is not
            normally described in the docstring but has been done here to act
            as an example.

        Warnings:
            Under no circumstances should `factor` be a torch.tensor entity.
            The `method` and `factor` parameters MUST be passed as positional
            arguments and NOT keyword arguments.

        """
        # Check that the method is of a known type
        if method not in _SymEigB.KNOWN_METHODS:
            raise ValueError('Unknown broadening method selected.')

        # Compute eigen-values & vectors using torch.symeig.
        w, v = torch.linalg.eigh(a)

        # Save tensors that will be needed in the backward pass
        ctx.save_for_backward(w, v)

        # Save the broadening factor and the selected broadening method.
        ctx.bf, ctx.bm = factor, method

        # Store dtype/device to prevent dtype/device mixing
        ctx.dtype, ctx.device = a.dtype, a.device

        # Return the eigenvalues and eigenvectors
        return w, v
//  [python]
def fit(self, data: DataLoader) -> EkfacInfluence:
        """
        Compute the KFAC blocks for each layer of the model, using the provided data.
        It then creates an EkfacRepresentation object that stores the KFAC blocks for
        each layer, their eigenvalue decomposition and diagonal values.
        """
        forward_x, grad_y = self._get_kfac_blocks(data)
        layers_evecs_a = {}
        layers_evect_g = {}
        layers_diags = {}
        for key in self.active_layers.keys():
            evals_a, evecs_a = torch.symeig(forward_x[key], eigenvectors=True)
            evals_g, evecs_g = torch.symeig(grad_y[key], eigenvectors=True)
            layers_evecs_a[key] = evecs_a
            layers_evect_g[key] = evecs_g
            layers_diags[key] = torch.kron(evals_g.view(-1, 1), evals_a.view(-1, 1))

        self.ekfac_representation = EkfacRepresentation(
            self.active_layers.keys(),
            self.active_layers.values(),
            layers_evecs_a.values(),
            layers_evect_g.values(),
            layers_diags.values(),
        )
        if self.update_diagonal:
            self._update_diag(data)
        return self
//  [python]
def fit(self, data: DataLoader) -> EkfacInfluence:
        """
        Compute the KFAC blocks for each layer of the model, using the provided data.
        It then creates an EkfacRepresentation object that stores the KFAC blocks for
        each layer, their eigenvalue decomposition and diagonal values.
        """
        forward_x, grad_y = self._get_kfac_blocks(data)
        layers_evecs_a = {}
        layers_evect_g = {}
        layers_diags = {}
        for key in self.active_layers.keys():
            evals_a, evecs_a = torch.linalg.eigh(forward_x[key])
            evals_g, evecs_g = torch.linalg.eigh(grad_y[key])
            layers_evecs_a[key] = evecs_a
            layers_evect_g[key] = evecs_g
            layers_diags[key] = torch.kron(evals_g.view(-1, 1), evals_a.view(-1, 1))

        self.ekfac_representation = EkfacRepresentation(
            self.active_layers.keys(),
            self.active_layers.values(),
            layers_evecs_a.values(),
            layers_evect_g.values(),
            layers_diags.values(),
        )
        if self.update_diagonal:
            self._update_diag(data)
        return self
//  [python]
def get_nystrom_features(K_mn, K_mm):
    s, Q = torch.eig(K_mm, eigenvectors=True)
    s = torch.real(s[:,0])
    D_inv_sqrt = torch.diag_embed(1.0/torch.sqrt(s))
    nystrom = torch.matmul(torch.transpose(Q, dim0=-2, dim1=-1), K_mn)
    nystrom = torch.matmul(D_inv_sqrt, nystrom)    
    return nystrom
//  [python]
def get_nystrom_features(K_mn, K_mm):
    s, Q = torch.linalg.eigh(K_mm)
    D_inv_sqrt = torch.diag_embed(1.0/torch.sqrt(s))
    nystrom = torch.matmul(torch.transpose(Q, dim0=-2, dim1=-1), K_mn)
    nystrom = torch.matmul(D_inv_sqrt, nystrom)    
    return nystrom
//  [python]
def matrix_operator(matrices: Tensor, operator) -> Tensor:
    """
    Matrix equivalent of an operator. Works batch-wise
    Porting of pyRiemann to pyTorch
    Original Author: Alexandre Barachant
    https://github.com/alexandrebarachant/pyRiemann
    """
    eigvals, eigvects = torch.eig(matrices, eigenvectors=True)
    eigvals = torch.diag_embed(operator(eigvals[...,0]))
    return eigvects @ eigvals @ eigvects.transpose(-2, -1)
//  [python]
def matrix_operator(matrices: Tensor, operator) -> Tensor:
    """
    Matrix equivalent of an operator. Works batch-wise
    Porting of pyRiemann to pyTorch
    Original Author: Alexandre Barachant
    https://github.com/alexandrebarachant/pyRiemann
    """
    eigvals, eigvects = torch.linalg.eigh(matrices, UPLO='L')
    eigvals = torch.diag_embed(operator(eigvals))
    return eigvects @ eigvals @ eigvects.transpose(-2, -1)
//  [python]
def generalized_eigh(A: Tensor, B: Tensor) -> tuple:
    r"""A workaround to solve a real symmetric generalized eigenvalue problem :math:`Av = \lambda Bv` using the eigenvalue decomposition of :math:`B^{-1/2}AB^{-1/2}`. This method is not numerically efficient.

    Parameters
    ----------
    A : Tensor

    B : Tensor

    Returns
    -------
    tuple
        Eigenvalues and eigenvectors of the generalized eigenvalue problem.
    """
    Lambda, Q = torch.symeig(B)
    rsqrt_Lambda = torch.diag(Lambda.rsqrt())
    rsqrt_B = Q @ rsqrt_Lambda
    _A = 0.5 * (
        rsqrt_B.T @ (A @ rsqrt_B) + rsqrt_B.T @ ((A.T) @ rsqrt_B)
    )  # Force Symmetrization
    values, _tmp_vecs = torch.symeig(_A)
    vectors = rsqrt_B @ _tmp_vecs
    return values, vectors
//  [python]
def generalized_eigh(A: Tensor, B: Tensor) -> tuple:
    r"""A workaround to solve a real symmetric generalized eigenvalue problem :math:`Av = \lambda Bv` using the eigenvalue decomposition of :math:`B^{-1/2}AB^{-1/2}`. This method is not numerically efficient.

    Parameters
    ----------
    A : Tensor

    B : Tensor

    Returns
    -------
    tuple
        Eigenvalues and eigenvectors of the generalized eigenvalue problem.
    """
    Lambda, Q = torch.linalg.eigh(B)
    rsqrt_Lambda = torch.diag(Lambda.rsqrt())
    rsqrt_B = Q @ rsqrt_Lambda
    _A = 0.5 * (
        rsqrt_B.T @ (A @ rsqrt_B) + rsqrt_B.T @ ((A.T) @ rsqrt_B)
    )  # Force Symmetrization
    values, _tmp_vecs = torch.linalg.eigh(_A)
    vectors = rsqrt_B @ _tmp_vecs
    return values, vectors
//  [python]
def map_positional_encoding(g, pos_enc_dim, use_unique_sign=True, use_unique_basis=True, use_eig_val=True):
    """
        Graph positional encoding v/ maximal axis projection
    """
    A = g.adjacency_matrix(scipy_fmt="csr").astype(np.double)
    N = sp.diags(dgl.backend.asnumpy(g.in_degrees()).clip(1) ** -0.5, dtype=np.double)
    A = torch.from_numpy((N * A * N).toarray())
    n, k = A.shape[0], pos_enc_dim
    E, U = torch.symeig(A, eigenvectors=True)
    E, U = E.round(decimals=14)[:-1], U[:, :-1]
    dim = min(n - 1, k)
    _, mult = torch.unique(E[-dim:], return_counts=True)
    ind = torch.cat([torch.LongTensor([0]), torch.cumsum(mult, dim=0)]) + max(n - 1 - k, 0)
    if use_unique_sign:
        for i in range(mult.shape[0]):
            if mult[i] == 1:
                U[:, ind[i]:ind[i + 1]] = unique_sign(U[:, ind[i]:ind[i + 1]])  # eliminate sign ambiguity
    if use_unique_basis:
        for i in range(mult.shape[0]):
            if mult[i] == 1:
                continue  # single eigenvector, no basis ambiguity
            try:
                U[:, ind[i]:ind[i + 1]] = unique_basis(U[:, ind[i]:ind[i + 1]])  # eliminate basis ambiguity
            except AssertionError:
                continue  # assumption violated, skip
    if use_eig_val:
        Lambda = torch.nn.ReLU()(torch.diag(E))
        U = U @ torch.sqrt(Lambda)
    if n - 1 < k:
        zeros = torch.zeros([n, k - n + 1])
        U = torch.cat([U, zeros], dim=-1)
    g.ndata['pos_enc'] = U[:, -k:]  # last k non-trivial eigenvectors
    return g
//  [python]
def map_positional_encoding(g, pos_enc_dim, use_unique_sign=True, use_unique_basis=True, use_eig_val=True):
    """
        Graph positional encoding v/ maximal axis projection
    """
    A = g.adjacency_matrix(scipy_fmt="csr").astype(np.double)
    N = sp.diags(dgl.backend.asnumpy(g.in_degrees()).clip(1) ** -0.5, dtype=np.double)
    A = torch.from_numpy((N * A * N).toarray())
    n, k = A.shape[0], pos_enc_dim
    E, U = torch.linalg.eigh(A)
    E, U = E.round(decimals=14)[:-1], U[:, :-1]
    dim = min(n - 1, k)
    _, mult = torch.unique(E[-dim:], return_counts=True)
    ind = torch.cat([torch.LongTensor([0]), torch.cumsum(mult, dim=0)]) + max(n - 1 - k, 0)
    if use_unique_sign:
        for i in range(mult.shape[0]):
            if mult[i] == 1:
                U[:, ind[i]:ind[i + 1]] = unique_sign(U[:, ind[i]:ind[i + 1]])  # eliminate sign ambiguity
    if use_unique_basis:
        for i in range(mult.shape[0]):
            if mult[i] == 1:
                continue  # single eigenvector, no basis ambiguity
            try:
                U[:, ind[i]:ind[i + 1]] = unique_basis(U[:, ind[i]:ind[i + 1]])  # eliminate basis ambiguity
            except AssertionError:
                continue  # assumption violated, skip
    if use_eig_val:
        Lambda = torch.nn.ReLU()(torch.diag(E))
        U = U @ torch.sqrt(Lambda)
    if n - 1 < k:
        zeros = torch.zeros([n, k - n + 1])
        U = torch.cat([U, zeros], dim=-1)
    g.ndata['pos_enc'] = U[:, -k:]  # last k non-trivial eigenvectors
    return g
//  [python]
def rho_fermi(
        self, H: torch.Tensor, mu: float
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Compute the equilibrium density matrix corresponding to `H`.
        Also return the energies and eigenvectors of `H`."""
        E, V = torch.symeig(H)
        f = fermi(E, mu, self.T)
        rho = torch.einsum("...ab, ...b, ...cb -> ...ac", V, f, V.conj())
        return self.packed_hermitian.pack(rho), E, V
//  [python]
def rho_fermi(
        self, H: torch.Tensor, mu: float
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Compute the equilibrium density matrix corresponding to `H`.
        Also return the energies and eigenvectors of `H`."""
        E, V = torch.linalg.eigh(H)
        f = fermi(E, mu, self.T)
        rho = torch.einsum("...ab, ...b, ...cb -> ...ac", V, f, V.conj())
        return self.packed_hermitian.pack(rho), E, V
//  [python]
def step(self):
        if self.weight_decay > 0:
            for p in self.model.parameters():
                p.grad.data.add_(self.weight_decay, p.data)

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

            la = self.damping + self.weight_decay

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

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

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

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

            v = v.view(p.grad.data.size())
            updates[p] = v

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

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

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

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

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

            la = self.damping + self.weight_decay

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

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

            if classname == 'Conv2d':
                p_grad_mat = p.grad.data.view(p.grad.data.size(0), -1)
            else:
                p_grad_mat = p.grad.data

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

            v = v.view(p.grad.data.size())
            updates[p] = v

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

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

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

        self.optim.step()
        self.steps += 1
//  [python]
def Fidelity(self, rho):
        """
        Quantum fidelity, ``rho`` is the calculated quantum state.

        When it is a pure state, calculated as defined, when it is a mixed state, 
        calculated as decomposed, the details can be found in ``Ultrafast quantum 
        state tomography with feed-forward neural networks``.
        """
        rho = rho / torch.trace(rho)
        Fq = torch.tensor(0)
        if self.ty_state == 'pure':
            tmp = torch.matmul(self.rho_star.T.conj(), rho)[0, 0]
            Fq = (tmp * tmp.conj()).real
        else:
            if self.pure_flag == 1:
                Fq = torch.matmul(torch.matmul(self.rho_p.T.conj(), rho), self.rho_p)[0, 0].real
            else:
                eigenvalues, eigenvecs = torch.eig(self.rho_star, eigenvectors=True)
                eigenvalues = torch.abs(eigenvalues[:,0])
                sqrt_rho = torch.matmul(eigenvecs * torch.sqrt(eigenvalues), eigenvecs.T.conj())  # sqrtm(self.rho_star)
                rho_tmp = torch.matmul(torch.matmul(sqrt_rho, rho), sqrt_rho)  # sqrtm(self.rho_star).dot(rho).dot(sqrtm(self.rho_star))

                try:
                    eigenvalues = torch.symeig(rho_tmp,eigenvectors=False)[0]  # fast, in some special cases, an error will be reported
                except Exception:
                    print('error')
                    eigenvalues = torch.eig(rho_tmp, eigenvectors=False)[:,0]  # low

                sqrt_eigvals = torch.sqrt(torch.abs(eigenvalues))
                Fq = torch.sum(sqrt_eigvals)**2  # trace(sqrtm(sqrtm(self.rho_star).dot(rho).dot(sqrtm(self.rho_star))))**2

        if Fq > 1:
            Fq = 1  # precision error
        return Fq
//  [python]
def Fidelity(self, rho):
        """
        Quantum fidelity, ``rho`` is the calculated quantum state.

        When it is a pure state, calculated as defined, when it is a mixed state, 
        calculated as decomposed, the details can be found in ``Ultrafast quantum 
        state tomography with feed-forward neural networks``.
        """
        rho = rho / torch.trace(rho)
        Fq = torch.tensor(0)
        if self.ty_state == 'pure':
            tmp = torch.matmul(self.rho_star.T.conj(), rho)[0, 0]
            Fq = (tmp * tmp.conj()).real
        else:
            if self.pure_flag == 1:
                Fq = torch.matmul(torch.matmul(self.rho_p.T.conj(), rho), self.rho_p)[0, 0].real
            else:
                eigenvalues, eigenvecs = torch.linalg.eigh(self.rho_star)
                eigenvalues = torch.abs(eigenvalues)
                sqrt_rho = torch.matmul(eigenvecs * torch.sqrt(eigenvalues), eigenvecs.T.conj())  # sqrtm(self.rho_star)
                rho_tmp = torch.matmul(torch.matmul(sqrt_rho, rho), sqrt_rho)  # sqrtm(self.rho_star).dot(rho).dot(sqrtm(self.rho_star))

                try:
                    eigenvalues = torch.linalg.eigvalsh(rho_tmp)  # fast, in some special cases, an error will be reported
                except Exception:
                    print('error')
                    eigenvalues = torch.linalg.eigvals(rho_tmp)  # low

                sqrt_eigvals = torch.sqrt(torch.abs(eigenvalues))
                Fq = torch.sum(sqrt_eigvals)**2  # trace(sqrtm(sqrtm(self.rho_star).dot(rho).dot(sqrtm(self.rho_star))))**2

        if Fq > 1:
            Fq = 1  # precision error
        return Fq
//  [python]
def rot_to_quat_pth(rot, unstack_inputs=False):
  if unstack_inputs:
    rot = [torch.moveaxis(x, -1, 0) for x in torch.moveaxis(rot, -2, 0)]
  [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot
  # pylint: disable=bad-whitespace
  k = [[ xx + yy + zz,      zy - yz,      xz - zx,      yx - xy,],
       [      zy - yz, xx - yy - zz,      xy + yx,      xz + zx,],
       [      xz - zx,      xy + yx, yy - xx - zz,      yz + zy,],
       [      yx - xy,      xz + zx,      yz + zy, zz - xx - yy,]]
  # pylint: enable=bad-whitespace
  k = (1./3.) * torch.stack([torch.stack(x, axis=-1) for x in k], axis=-2)
  # Get eigenvalues in non-decreasing order and associated.
  _, qs = torch.symeig(k, eigenvectors=True)
  return qs[..., -1]
//  [python]
def rot_to_quat_pth(rot, unstack_inputs=False):
  if unstack_inputs:
    rot = [torch.moveaxis(x, -1, 0) for x in torch.moveaxis(rot, -2, 0)]
  [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot
  # pylint: disable=bad-whitespace
  k = [[ xx + yy + zz,      zy - yz,      xz - zx,      yx - xy,],
       [      zy - yz, xx - yy - zz,      xy + yx,      xz + zx,],
       [      xz - zx,      xy + yx, yy - xx - zz,      yz + zy,],
       [      yx - xy,      xz + zx,      yz + zy, zz - xx - yy,]]
  # pylint: enable=bad-whitespace
  k = (1./3.) * torch.stack([torch.stack(x, axis=-1) for x in k], axis=-2)
  # Get eigenvalues in non-decreasing order and associated.
  _, qs = torch.linalg.eigh(k)
  return qs[..., -1]
//  [python]
def sym_matrix_pow(matrix: torch.Tensor, p: float) -> torch.Tensor:
    r"""
    Power of a matrix using Eigen Decomposition.
    Args:
        matrix: a batch of matrices
        p: power
    Returns:
        Power of a matrix
    """
    vals, vecs = torch.symeig(matrix)
    vals[vals > 0] = vals[vals > 0].pow(p)
    matrix_pow = vecs @ torch.diag(vals) @ vecs.T
    return matrix_pow
//  [python]
def sym_matrix_pow(matrix: torch.Tensor, p: float) -> torch.Tensor:
    r"""
    Power of a matrix using Eigen Decomposition.
    Args:
        matrix: a batch of matrices
        p: power
    Returns:
        Power of a matrix
    """
    vals, vecs = torch.linalg.eigh(matrix)
    vals[vals > 0] = vals[vals > 0].pow(p)
    matrix_pow = vecs @ torch.diag(vals) @ vecs.T
    return matrix_pow
//  [python]
def sorted_eig(cov):  # For now we sort 'by convention'. For PCA the sorting is key.
    lambdas, vs = torch.eig(cov, eigenvectors=True)   # cov is symmetric so we can use eigh instead of eig
    lambdas = lambdas[:,0]
    lambdas, indices = torch.sort(lambdas, dim=0, descending=True)
    vs = torch.index_select(vs, 0, indices)
    return lambdas, vs
//  [python]
def sorted_eig(cov):  # For now we sort 'by convention'. For PCA the sorting is key.
    lambdas, vs = torch.linalg.eigh(cov)   # cov is symmetric so we can use eigh instead of eig
    lambdas, indices = torch.sort(lambdas, dim=0, descending=True)
    vs = torch.index_select(vs, 0, indices)
    return lambdas, vs
//  [python]
def _log_prob(self, images_flat):
        L, Q = torch.eig(self.sigma, eigenvectors=True)
        half_inv_sigma = Q @ torch.diag_embed(1 / torch.sqrt(L[:, 0]).clamp(min=self.eps)) @ Q.transpose(-1, -2)
        norm_images = images_flat - self.mu.unsqueeze(-1)
        return - 0.5 * self.num_subjects * torch.log(torch.as_tensor(torch.pi * 2,
                                                                     dtype=self.dtype, device=self.device)) \
               - 0.5 * torch.prod(L[:, 0], dim=-1, keepdim=True).clamp(min=self.eps).log() \
               - 0.5 * torch.square(half_inv_sigma @ norm_images).sum(dim=2)
//  [python]
def _log_prob(self, images_flat):
        L, Q = torch.linalg.eigh(self.sigma)
        half_inv_sigma = Q @ torch.diag_embed(1 / torch.sqrt(L).clamp(min=self.eps)) @ Q.transpose(-1, -2)
        norm_images = images_flat - self.mu.unsqueeze(-1)
        return - 0.5 * self.num_subjects * torch.log(torch.as_tensor(torch.pi * 2,
                                                                     dtype=self.dtype, device=self.device)) \
               - 0.5 * torch.prod(L, dim=-1, keepdim=True).clamp(min=self.eps).log() \
               - 0.5 * torch.square(half_inv_sigma @ norm_images).sum(dim=2)
//  [python]
def solve(self, x):
        m = x.shape[1]

        if True:  # self.enable_symmetric:
            l, u = torch.symeig(x)
        else:
            l, u = torch.eig(x)

        u = self.uniform_solution_direction(u)

        # The greatest Eigenvalue first, the default is the greatest last
        l = torch.flip(l, dims=[1])
        u = torch.flip(u, dims=[2])

        if (self.num_eigen_values <= 0)  or (self.num_eigen_values > m):
            self.num_eigen_values = m

        return l[:, :self.num_eigen_values], u[:, :, :self.num_eigen_values]
//  [python]
def solve(self, x):
        m = x.shape[1]

        if True:  # self.enable_symmetric:
            l, u = torch.linalg.eigh(x)
        else:
            l, u = torch.linalg.eig(x)

        u = self.uniform_solution_direction(u)

        # The greatest Eigenvalue first, the default is the greatest last
        l = torch.flip(l, dims=[1])
        u = torch.flip(u, dims=[2])

        if (self.num_eigen_values <= 0)  or (self.num_eigen_values > m):
            self.num_eigen_values = m

        return l[:, :self.num_eigen_values], u[:, :, :self.num_eigen_values]