jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def fit_Nystrom(Q:Tensor, y:Tensor, train_mask:Tensor, nugget:Union[float, Tensor]=1e-2):
    """
    Make predictions for a range of nugget values in one shot.
    For classification problems, the prediction target is classification probability.
    For regression problems, the prediction target is output value.

    Args:
        Q (Tensor[float]): square root of the GNNGP Kernel matrix.
        y (Tensor[float] or Tensor[int]): prediction target.
            For int type, the task is a classification problem.
            For float type, the task is a regression problem.
        train_mask (Tensor[bool]): mask for training nodes.
        nugget (float or Tensor[float]): the nugget used in posterior inference.
            (default: 1e-2 x mean diagonal element of QQ^T)
    """
    mb = train_mask
    if isinstance(nugget, float): nugget = torch.tensor([nugget])
    if y.dtype in [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64]:
        yb = torch.nn.functional.one_hot(y[mb]).to(torch.float)
    else: yb = y[mb]
    fitted = torch.zeros((len(nugget), Q.shape[0], *yb.shape[1:]), device=Q.device)
    w, v = torch.symeig(Q[mb].T @ Q[mb])
    d = torch.mean(torch.sum(Q**2, 1))
    shape = (-1, 1) if len(yb.shape) >= 2 else (-1,)
    temp = v.T @ (Q[mb].T @ yb)
    for i, eps in enumerate(nugget):
        fitted[i] = Q @ (v @ (temp/(w+eps*d).view(shape)))
    return fitted[0] if len(nugget) == 1 else fitted
//  [python]
def fit_Nystrom(Q:Tensor, y:Tensor, train_mask:Tensor, nugget:Union[float, Tensor]=1e-2):
    """
    Make predictions for a range of nugget values in one shot.
    For classification problems, the prediction target is classification probability.
    For regression problems, the prediction target is output value.

    Args:
        Q (Tensor[float]): square root of the GNNGP Kernel matrix.
        y (Tensor[float] or Tensor[int]): prediction target.
            For int type, the task is a classification problem.
            For float type, the task is a regression problem.
        train_mask (Tensor[bool]): mask for training nodes.
        nugget (float or Tensor[float]): the nugget used in posterior inference.
            (default: 1e-2 x mean diagonal element of QQ^T)
    """
    mb = train_mask
    if isinstance(nugget, float): nugget = torch.tensor([nugget])
    if y.dtype in [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64]:
        yb = torch.nn.functional.one_hot(y[mb]).to(torch.float)
    else: yb = y[mb]
    fitted = torch.zeros((len(nugget), Q.shape[0], *yb.shape[1:]), device=Q.device)
    w, v = torch.linalg.eigh(Q[mb].T @ Q[mb])
    d = torch.mean(torch.sum(Q**2, 1))
    shape = (-1, 1) if len(yb.shape) >= 2 else (-1,)
    temp = v.T @ (Q[mb].T @ yb)
    for i, eps in enumerate(nugget):
        fitted[i] = Q @ (v @ (temp/(w+eps*d).view(shape)))
    return fitted[0] if len(nugget) == 1 else fitted
//  [python]
def get_gaussian_ellipsiod(
    mean: TensorType[3],
    cov: TensorType[3, 3],
    n_std: int = 2,
    color="lightblue",
    opacity: float = 0.5,
    resolution: int = 20,
    name: str = "ellipse",
) -> go.Mesh3d:  # type: ignore
    """Get a plotly ellipsoid for a Gaussian.

    Args:
        mean: mean of the Gaussian.
        cov: covariance of the Gaussian.
        n_std: Standard deviation to visualize. Defaults to 2 (95% confidence).
        color: Color of the ellipsoid. Defaults to None.
        opacity: Opacity of the ellipsoid. Defaults to 0.5.
        resolution: Resolution of the ellipsoid. Defaults to 20.
        name: Name of the ellipsoid. Defaults to "ellipse".

    Returns:
        ellipsoid object.
    """

    phi = torch.linspace(0, 2 * torch.pi, resolution)
    theta = torch.linspace(-torch.pi / 2, torch.pi / 2, resolution)
    phi, theta = torch.meshgrid(phi, theta, indexing="ij")

    x = torch.cos(theta) * torch.sin(phi)
    y = torch.cos(theta) * torch.cos(phi)
    z = torch.sin(theta)
    pts = torch.stack((x, y, z), dim=-1)

    eigenvals, eigenvecs = torch.symeig(cov, eigenvectors=True)
    idx = torch.sum(cov, dim=0).argsort()
    idx = eigenvals[idx].argsort()
    eigenvals = eigenvals[idx]
    eigenvecs = eigenvecs[:, idx]

    scaling = torch.sqrt(eigenvals) * n_std
    pts = pts * scaling

    pts = pts @ eigenvecs.t()

    pts += mean

    return go.Mesh3d(
        {
            "x": pts[:, :, 0].flatten(),
            "y": pts[:, :, 1].flatten(),
            "z": pts[:, :, 2].flatten(),
            "alphahull": 0,
            "opacity": opacity,
            "color": color,
            "name": name,
        }
    )
//  [python]
def get_gaussian_ellipsiod(
    mean: TensorType[3],
    cov: TensorType[3, 3],
    n_std: int = 2,
    color="lightblue",
    opacity: float = 0.5,
    resolution: int = 20,
    name: str = "ellipse",
) -> go.Mesh3d:  # type: ignore
    """Get a plotly ellipsoid for a Gaussian.

    Args:
        mean: mean of the Gaussian.
        cov: covariance of the Gaussian.
        n_std: Standard deviation to visualize. Defaults to 2 (95% confidence).
        color: Color of the ellipsoid. Defaults to None.
        opacity: Opacity of the ellipsoid. Defaults to 0.5.
        resolution: Resolution of the ellipsoid. Defaults to 20.
        name: Name of the ellipsoid. Defaults to "ellipse".

    Returns:
        ellipsoid object.
    """

    phi = torch.linspace(0, 2 * torch.pi, resolution)
    theta = torch.linspace(-torch.pi / 2, torch.pi / 2, resolution)
    phi, theta = torch.meshgrid(phi, theta, indexing="ij")

    x = torch.cos(theta) * torch.sin(phi)
    y = torch.cos(theta) * torch.cos(phi)
    z = torch.sin(theta)
    pts = torch.stack((x, y, z), dim=-1)

    eigenvals, eigenvecs = torch.linalg.eigh(cov)
    idx = torch.sum(cov, dim=0).argsort()
    idx = eigenvals[idx].argsort()
    eigenvals = eigenvals[idx]
    eigenvecs = eigenvecs[:, idx]

    scaling = torch.sqrt(eigenvals) * n_std
    pts = pts * scaling

    pts = pts @ eigenvecs.t()

    pts += mean

    return go.Mesh3d(
        {
            "x": pts[:, :, 0].flatten(),
            "y": pts[:, :, 1].flatten(),
            "z": pts[:, :, 2].flatten(),
            "alphahull": 0,
            "opacity": opacity,
            "color": color,
            "name": name,
        }
    )
//  [python]
def forward(self,previousSample,energy_function,numSteps=1,stepSize=0.01,k=5,**kwargs):
        new_draw = previousSample.clone().detach().requires_grad_(True)
        #prev_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
        energy,hess = energy_function(kwargs["decoder"],new_draw)

        if isinstance(energy,str):
            if energy=="nan":
                return "eig failed"
        #print(check_symmetric(hess))
        #norm_gradient = gradient/torch.norm(gradient)

        #norm_gradient = gradient/torch.norm(gradient)
        self.t += 0.1
        for _ in range(1):
            #U,S,V = torch.linalg.svd(hess)
            #Mntm = torch.randn_like(new_draw)
            #w = Mntm - torch.mm(Mntm,torch.unsqueeze(norm_gradient,-1))*norm_gradient

            decomposition = torch.symeig(hess)
            S = decomposition.eigenvalues
            V = decomposition.eigenvectors
            S,indices = torch.sort(S,descending=True)
            V = V[:,indices]
            num_components = 1
            Vk = V[:,:num_components]
            Sk = S[:num_components]
            beta = torch.randn(num_components)
            beta = beta/torch.norm(beta)
            direction = torch.zeros(new_draw.size()).float().cuda()

            for compIndex in range(num_components):
                #if Sk[compIndex]>0:
                direction += self.t*Vk[:,compIndex]
                #else:
                #    direction += beta[compIndex]*(-1*Vk[:,compIndex])

            #norm_direction = torch.norm(direction)
            #direction = direction/norm_direction
            '''

            delta = 1e-3
            betasq_lambda = 0.0
            for compIndex in range(num_components):
                betasq_lambda += ((beta[compIndex]/norm_direction)**2 * torch.abs(Sk[compIndex]))
               
            stepSize = torch.sqrt((2.0*delta)/betasq_lambda)
            stepSize = max(0.2,min(0.7,stepSize))
            #print(stepSize)
            stepSize = 0.1
            #stepSize = 0.1
            '''
            stepSize = 1.0

            new_draw = (new_draw+stepSize*direction).clone().detach().requires_grad_(True)
            #new_draw = direction.clone().detach().requires_grad_(True)
            #current_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
            #print("Energy:",current_U)
            #norm_gradient = gradient/torch.norm(gradient)
        draw_final = new_draw.clone().detach()

        return draw_final
//  [python]
def forward(self,previousSample,energy_function,numSteps=1,stepSize=0.01,k=5,**kwargs):
        new_draw = previousSample.clone().detach().requires_grad_(True)
        #prev_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
        energy,hess = energy_function(kwargs["decoder"],new_draw)

        if isinstance(energy,str):
            if energy=="nan":
                return "eig failed"
        #print(check_symmetric(hess))
        #norm_gradient = gradient/torch.norm(gradient)

        #norm_gradient = gradient/torch.norm(gradient)
        self.t += 0.1
        for _ in range(1):
            #U,S,V = torch.linalg.svd(hess)
            #Mntm = torch.randn_like(new_draw)
            #w = Mntm - torch.mm(Mntm,torch.unsqueeze(norm_gradient,-1))*norm_gradient

            decomposition = torch.linalg.eigh(hess)
            S = decomposition.eigenvalues
            V = decomposition.eigenvectors
            S,indices = torch.sort(S,descending=True)
            V = V[:,indices]
            num_components = 1
            Vk = V[:,:num_components]
            Sk = S[:num_components]
            beta = torch.randn(num_components)
            beta = beta/torch.norm(beta)
            direction = torch.zeros(new_draw.size()).float().cuda()

            for compIndex in range(num_components):
                #if Sk[compIndex]>0:
                direction += self.t*Vk[:,compIndex]
                #else:
                #    direction += beta[compIndex]*(-1*Vk[:,compIndex])

            #norm_direction = torch.norm(direction)
            #direction = direction/norm_direction
            '''

            delta = 1e-3
            betasq_lambda = 0.0
            for compIndex in range(num_components):
                betasq_lambda += ((beta[compIndex]/norm_direction)**2 * torch.abs(Sk[compIndex]))
               
            stepSize = torch.sqrt((2.0*delta)/betasq_lambda)
            stepSize = max(0.2,min(0.7,stepSize))
            #print(stepSize)
            stepSize = 0.1
            #stepSize = 0.1
            '''
            stepSize = 1.0

            new_draw = (new_draw+stepSize*direction).clone().detach().requires_grad_(True)
            #new_draw = direction.clone().detach().requires_grad_(True)
            #current_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
            #print("Energy:",current_U)
            #norm_gradient = gradient/torch.norm(gradient)
        draw_final = new_draw.clone().detach()

        return draw_final
//  [python]
def init_spatial_model(self):
        self.start_idx = 0
        self.Q_FMM = torch.tile(
            torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device), [self.n_freq, 1, 1]
        )
        self.G_NFM = self.g_eps * torch.ones(
            [self.n_source, self.n_freq, self.n_mic], dtype=self.TYPE_FLOAT, device=self.device
        )
        for m in range(self.n_mic):
            self.G_NFM[m % self.n_source, :, m] = 1

        if "circular" in self.init_SCM:
            pass
        elif "obs" in self.init_SCM:
            if hasattr(self, "XX_FTMM"):
                XX_FMM = self.XX_FTMM.sum(axis=1)
            else:
                XX_FMM = torch.einsum("fti, ftj -> fij", self.X_FTM, self.X_FTM.conj())
            _, eig_vec_FMM = torch.symeig(XX_FMM, eigenvectors=True)
            eig_vec_FMM = eig_vec_FMM[:, :, range(self.n_mic - 1, -1, -1)]
            self.Q_FMM = eig_vec_FMM.permute(0, 2, 1).conj()
        elif "twostep" == self.init_SCM:
            if self.n_iter_init >= self.n_iter:
                print(
                    "\n------------------------------------------------------------------\n"
                    f"Warning: n_iter_init must be smaller than n_iter (= {self.n_iter}).\n"
                    f"n_iter_init is changed from {self.n_iter_init} to {self.n_iter // 3}"
                    "\n------------------------------------------------------------------\n"
                )
                self.n_iter_init = self.n_iter // 3

            self.start_idx = self.n_iter_init
            separater_init = FastMNMF1(
                n_source=self.n_source,
                n_basis=2,
                device=self.device,
                init_SCM="circular",
                n_bit=self.n_bit,
                g_eps=self.g_eps,
            )
            separater_init.load_spectrogram(self.X_FTM, self.sample_rate)
            separater_init.solve(n_iter=self.start_idx, save_wav=False)

            self.Q_FMM = separater_init.Q_FMM
            self.G_NFM = separater_init.G_NFM
        else:
            raise ValueError("init_SCM should be circular, obs, or twostep.")

        self.G_NFM /= self.G_NFM.sum(axis=2)[..., None]
        self.normalize()
//  [python]
def init_spatial_model(self):
        self.start_idx = 0
        self.Q_FMM = torch.tile(
            torch.eye(self.n_mic, dtype=self.TYPE_COMPLEX, device=self.device), [self.n_freq, 1, 1]
        )
        self.G_NFM = self.g_eps * torch.ones(
            [self.n_source, self.n_freq, self.n_mic], dtype=self.TYPE_FLOAT, device=self.device
        )
        for m in range(self.n_mic):
            self.G_NFM[m % self.n_source, :, m] = 1

        if "circular" in self.init_SCM:
            pass
        elif "obs" in self.init_SCM:
            if hasattr(self, "XX_FTMM"):
                XX_FMM = self.XX_FTMM.sum(axis=1)
            else:
                XX_FMM = torch.einsum("fti, ftj -> fij", self.X_FTM, self.X_FTM.conj())
            _, eig_vec_FMM = torch.linalg.eigh(XX_FMM)
            eig_vec_FMM = eig_vec_FMM[:, :, range(self.n_mic - 1, -1, -1)]
            self.Q_FMM = eig_vec_FMM.permute(0, 2, 1).conj()
        elif "twostep" == self.init_SCM:
            if self.n_iter_init >= self.n_iter:
                print(
                    "\n------------------------------------------------------------------\n"
                    f"Warning: n_iter_init must be smaller than n_iter (= {self.n_iter}).\n"
                    f"n_iter_init is changed from {self.n_iter_init} to {self.n_iter // 3}"
                    "\n------------------------------------------------------------------\n"
                )
                self.n_iter_init = self.n_iter // 3

            self.start_idx = self.n_iter_init
            separater_init = FastMNMF1(
                n_source=self.n_source,
                n_basis=2,
                device=self.device,
                init_SCM="circular",
                n_bit=self.n_bit,
                g_eps=self.g_eps,
            )
            separater_init.load_spectrogram(self.X_FTM, self.sample_rate)
            separater_init.solve(n_iter=self.start_idx, save_wav=False)

            self.Q_FMM = separater_init.Q_FMM
            self.G_NFM = separater_init.G_NFM
        else:
            raise ValueError("init_SCM should be circular, obs, or twostep.")

        self.G_NFM /= self.G_NFM.sum(axis=2)[..., None]
        self.normalize()
//  [python]
def generalized_eigenvalue_decomposition(a: torch.Tensor, b: torch.Tensor, eps=1e-6):
    """Solves the generalized eigenvalue decomposition through Cholesky decomposition.

    ported from https://github.com/asteroid-team/asteroid/blob/master/asteroid/dsp/beamforming.py#L464

    a @ e_vec = e_val * b @ e_vec
    |
    |   Cholesky decomposition on `b`:
    |       b = L @ L^H, where `L` is a lower triangular matrix
    |
    |   Let C = L^-1 @ a @ L^-H, it is Hermitian.
    |
    => C @ y = lambda * y
    => e_vec = L^-H @ y

    Reference: https://www.netlib.org/lapack/lug/node54.html

    Args:
        a: A complex Hermitian or real symmetric matrix whose eigenvalues and
            eigenvectors will be computed. (..., C, C)
        b: A complex Hermitian or real symmetric definite positive matrix. (..., C, C)
    Returns:
        e_val: generalized eigenvalues (ascending order)
        e_vec: generalized eigenvectors
    """  # noqa: H405, E501
    try:
        cholesky = torch.linalg.cholesky(b)
    except RuntimeError:
        b = tik_reg(b, reg=eps, eps=eps)
        cholesky = torch.linalg.cholesky(b)
    inv_cholesky = cholesky.inverse()
    # Compute C matrix L⁻1 a L^-H
    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: torch.Tensor, b: torch.Tensor, eps=1e-6):
    """Solves the generalized eigenvalue decomposition through Cholesky decomposition.

    ported from https://github.com/asteroid-team/asteroid/blob/master/asteroid/dsp/beamforming.py#L464

    a @ e_vec = e_val * b @ e_vec
    |
    |   Cholesky decomposition on `b`:
    |       b = L @ L^H, where `L` is a lower triangular matrix
    |
    |   Let C = L^-1 @ a @ L^-H, it is Hermitian.
    |
    => C @ y = lambda * y
    => e_vec = L^-H @ y

    Reference: https://www.netlib.org/lapack/lug/node54.html

    Args:
        a: A complex Hermitian or real symmetric matrix whose eigenvalues and
            eigenvectors will be computed. (..., C, C)
        b: A complex Hermitian or real symmetric definite positive matrix. (..., C, C)
    Returns:
        e_val: generalized eigenvalues (ascending order)
        e_vec: generalized eigenvectors
    """  # noqa: H405, E501
    try:
        cholesky = torch.linalg.cholesky(b)
    except RuntimeError:
        b = tik_reg(b, reg=eps, eps=eps)
        cholesky = torch.linalg.cholesky(b)
    inv_cholesky = cholesky.inverse()
    # Compute C matrix L⁻1 a L^-H
    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 eig_fun(hmat, fun):
    """Apply a function to the eigenvalues of a symmetric (hermitian) matrix
    """
    e, u = torch.symeig(hmat)
    return u @ torch.diag(fun(e)) @ u.T
//  [python]
def eig_fun(hmat, fun):
    """Apply a function to the eigenvalues of a symmetric (hermitian) matrix
    """
    e, u = torch.linalg.eigh(hmat)
    return u @ torch.diag(fun(e)) @ u.T
//  [python]
def sqrtmh(A: torch.Tensor):
    # Credits to
    """Compute the square root of a Symmetric or Hermitian positive definite matrix or batch of matrices. Credits to  `https://github.com/pytorch/pytorch/issues/25481#issuecomment-1032789228 <https://github.com/pytorch/pytorch/issues/25481#issuecomment-1032789228>`_."""
    L, Q = torch.eig(A)
    zero = torch.zeros((), device=L.device, dtype=L.dtype)
    threshold = L.max(-1).values * L.size(-1) * torch.finfo(L.dtype).eps
    L = L.where(L > threshold.unsqueeze(-1), zero)  # zero out small components
    return (Q * L.sqrt().unsqueeze(-2)) @ Q.mH
//  [python]
def sqrtmh(A: torch.Tensor):
    # Credits to
    """Compute the square root of a Symmetric or Hermitian positive definite matrix or batch of matrices. Credits to  `https://github.com/pytorch/pytorch/issues/25481#issuecomment-1032789228 <https://github.com/pytorch/pytorch/issues/25481#issuecomment-1032789228>`_."""
    L, Q = torch.linalg.eigh(A)
    zero = torch.zeros((), device=L.device, dtype=L.dtype)
    threshold = L.max(-1).values * L.size(-1) * torch.finfo(L.dtype).eps
    L = L.where(L > threshold.unsqueeze(-1), zero)  # zero out small components
    return (Q * L.sqrt().unsqueeze(-2)) @ Q.mH
//  [python]
def get_eigenvectors(factors: Dict[Module, Tensor]) -> Dict[Module, Tensor]:
    """Computes the eigenvectors of KFAC factors.

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

    Returns:
        A dict mapping layers to lists containing the first and second KFAC factors eigenvectors.
    """
    eigenvectors = dict()
    for layer, (xxt, ggt) in factors.items():
        sym_xxt, sym_ggt = xxt + xxt.t(), ggt + ggt.t()
        _, xxt_eigvecs = torch.linalg.eigh(sym_xxt, UPLO='U')
        _, ggt_eigvecs = torch.linalg.eigh(sym_ggt, UPLO='U')
        eigenvectors[layer] = (xxt_eigvecs, ggt_eigvecs)
    return eigenvectors
//  [python]
def rot_to_quat(
    rot: torch.Tensor,
):
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot

    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,
        ],
    ]

    k = (1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

    _, vectors = torch.symeig(k, eigenvectors=True)
    return vectors[..., -1]
//  [python]
def rot_to_quat(
    rot: torch.Tensor,
):
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot

    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,
        ],
    ]

    k = (1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

    _, vectors = torch.linalg.eigh(k)
    return vectors[..., -1]
//  [python]
def fit(self, data):

        assert len(data) == 2

        c_tilde = estimate_c_tilde_matrix(data[0], data[1])
        self._eigenvalues, _ = torch.eig(c_tilde)
        self._score = torch.sum(self._eigenvalues**2) + 1
        self._is_fitted = True

        return self
//  [python]
def fit(self, data):

        assert len(data) == 2

        c_tilde = estimate_c_tilde_matrix(data[0], data[1])
        self._eigenvalues, _ = torch.linalg.eigh(c_tilde)
        self._score = torch.sum(self._eigenvalues**2) + 1
        self._is_fitted = True

        return self
//  [python]
def norm_sender_normalized(pyr, num_scale=2, num_bands=6, blksz=3, eps=1e-12):
    r"""Normalize pyramid with local spatial neighbor and band neighbor
    """
    border = blksz // 2
    guardband = 16
    subbands = []
    for si in range(num_scale):
        for bi in range(num_bands):
            idx = si * num_bands + bi
            current_band = pyr[idx]

            N = blksz**2

            # 3x3 window pixels
            tmp = F.unfold(current_band.unsqueeze(1), 3, stride=1)
            tmp = tmp.transpose(1, 2)
            b, hw = tmp.shape[:2]
            # parent pixels
            parent_idx = idx + num_bands
            if parent_idx < len(pyr):
                tmp_parent = pyr[parent_idx]
                tmp_parent = imresize(tmp_parent, sizes=current_band.shape[-2:])
                tmp_parent = tmp_parent[:, border:-border, border:-border].reshape(b, hw, 1)
                tmp = torch.cat((tmp, tmp_parent), dim=-1)
                N += 1
            # neighbor band pixels
            for ni in range(num_bands):
                if ni != bi:
                    ni_idx = si * num_bands + ni
                    tmp_nei = pyr[ni_idx]
                    tmp_nei = tmp_nei[:, border:-border, border:-border].reshape(b, hw, 1)
                    tmp = torch.cat((tmp, tmp_nei), dim=-1)
            C_x = tmp.transpose(1, 2) @ tmp / tmp.shape[1]
            # correct possible negative eigenvalue
            L, Q = torch.symeig(C_x)
            L_pos = L * (L > 0)
            L_pos_sum = L_pos.sum(dim=1, keepdim=True)
            L = L_pos * L.sum(dim=1, keepdim=True) / (L_pos_sum + (L_pos_sum == 0).to(L.dtype))
            C_x = Q @ torch.diag_embed(L) @ Q.transpose(1, 2)

            o_c = current_band[:, border:-border, border:-border]
            b, h, w = o_c.shape
            o_c = o_c.reshape(b, hw)
            o_c = o_c - o_c.mean(dim=1, keepdim=True)

            if tmp.shape[1] >= 2e5: # To avoid out of GPU memory
                C_x = C_x.cpu()
                tmp = tmp.cpu()
            if hasattr(torch.linalg, 'lstsq'):
                tmp_y = torch.linalg.lstsq(C_x.transpose(1, 2), tmp.transpose(1, 2)).solution.transpose(1, 2) * tmp / N
            else:
                warn(
                    "For numerical stability, we use torch.linal.lstsq to calculate matrix inverse for PyTorch > 1.9.0. The results might be slightly different if you use older version of PyTorch.")
                tmp_y = (tmp @ torch.linalg.pinv(C_x)) * tmp / N
            tmp_y = tmp_y.to(o_c)

            z = tmp_y.sum(dim=2).sqrt()
            mask = z != 0
            g_c = o_c * mask / (z * mask + eps)
            g_c = g_c.reshape(b, h, w)

            gb = int(guardband / (2**(si)))
            g_c = g_c[:, gb:-gb, gb:-gb]
            g_c = g_c - g_c.mean(dim=(1, 2), keepdim=True)
            subbands.append(g_c)

    return subbands
//  [python]
def norm_sender_normalized(pyr, num_scale=2, num_bands=6, blksz=3, eps=1e-12):
    r"""Normalize pyramid with local spatial neighbor and band neighbor
    """
    border = blksz // 2
    guardband = 16
    subbands = []
    for si in range(num_scale):
        for bi in range(num_bands):
            idx = si * num_bands + bi
            current_band = pyr[idx]

            N = blksz**2

            # 3x3 window pixels
            tmp = F.unfold(current_band.unsqueeze(1), 3, stride=1)
            tmp = tmp.transpose(1, 2)
            b, hw = tmp.shape[:2]
            # parent pixels
            parent_idx = idx + num_bands
            if parent_idx < len(pyr):
                tmp_parent = pyr[parent_idx]
                tmp_parent = imresize(tmp_parent, sizes=current_band.shape[-2:])
                tmp_parent = tmp_parent[:, border:-border, border:-border].reshape(b, hw, 1)
                tmp = torch.cat((tmp, tmp_parent), dim=-1)
                N += 1
            # neighbor band pixels
            for ni in range(num_bands):
                if ni != bi:
                    ni_idx = si * num_bands + ni
                    tmp_nei = pyr[ni_idx]
                    tmp_nei = tmp_nei[:, border:-border, border:-border].reshape(b, hw, 1)
                    tmp = torch.cat((tmp, tmp_nei), dim=-1)
            C_x = tmp.transpose(1, 2) @ tmp / tmp.shape[1]
            # correct possible negative eigenvalue
            L, Q = torch.linalg.eigh(C_x)
            L_pos = L * (L > 0)
            L_pos_sum = L_pos.sum(dim=1, keepdim=True)
            L = L_pos * L.sum(dim=1, keepdim=True) / (L_pos_sum + (L_pos_sum == 0).to(L.dtype))
            C_x = Q @ torch.diag_embed(L) @ Q.transpose(1, 2)

            o_c = current_band[:, border:-border, border:-border]
            b, h, w = o_c.shape
            o_c = o_c.reshape(b, hw)
            o_c = o_c - o_c.mean(dim=1, keepdim=True)

            if tmp.shape[1] >= 2e5: # To avoid out of GPU memory
                C_x = C_x.cpu()
                tmp = tmp.cpu()
            if hasattr(torch.linalg, 'lstsq'):
                tmp_y = torch.linalg.lstsq(C_x.transpose(1, 2), tmp.transpose(1, 2)).solution.transpose(1, 2) * tmp / N
            else:
                warn(
                    "For numerical stability, we use torch.linal.lstsq to calculate matrix inverse for PyTorch > 1.9.0. The results might be slightly different if you use older version of PyTorch.")
                tmp_y = (tmp @ torch.linalg.pinv(C_x)) * tmp / N
            tmp_y = tmp_y.to(o_c)

            z = tmp_y.sum(dim=2).sqrt()
            mask = z != 0
            g_c = o_c * mask / (z * mask + eps)
            g_c = g_c.reshape(b, h, w)

            gb = int(guardband / (2**(si)))
            g_c = g_c[:, gb:-gb, gb:-gb]
            g_c = g_c - g_c.mean(dim=(1, 2), keepdim=True)
            subbands.append(g_c)

    return subbands
//  [python]
def step(self, direction: MatrixArray, step_size: float) -> None:
        el = self.system.electrons
        # Move auxiliary Hamiltonian accounting for accumulated rotations:
        H_aux = el.eig.diag_embed() + step_size * (  # current Hamiltonian
            dagger(self._rot_prev) @ (direction.M @ self._rot_prev)
        )
        # Update rotations to re-diagonalize auxiliary Hamiltonian
        el.eig, V = torch.symeig(H_aux)
        self._rot_prev = self._rot_prev @ V
        el.C = el.C @ V
//  [python]
def step(self, direction: MatrixArray, step_size: float) -> None:
        el = self.system.electrons
        # Move auxiliary Hamiltonian accounting for accumulated rotations:
        H_aux = el.eig.diag_embed() + step_size * (  # current Hamiltonian
            dagger(self._rot_prev) @ (direction.M @ self._rot_prev)
        )
        # Update rotations to re-diagonalize auxiliary Hamiltonian
        el.eig, V = torch.linalg.eigh(H_aux)
        self._rot_prev = self._rot_prev @ V
        el.C = el.C @ V
//  [python]
def spectral_norm(k_hat, k):
    difference = k_hat - k
    evals, _ = torch.symeig(difference, eigenvectors=False)
    lambda_min = evals[0]
    lambda_max = evals[-1]

    assert lambda_max >= lambda_min
    spectral_norm = max(lambda_max, -lambda_min)

    evals, _ = torch.symeig(k, eigenvectors=False)
    lambda_min = evals[0]
    lambda_max = evals[-1]

    assert lambda_max >= lambda_min
    ref_spectral_norm = max(lambda_max, -lambda_min)

    rel_spectral_norm = spectral_norm / ref_spectral_norm

    return spectral_norm, rel_spectral_norm
//  [python]
def spectral_norm(k_hat, k):
    difference = k_hat - k
    evals, _ = torch.linalg.eigh(difference)
    lambda_min = evals[0]
    lambda_max = evals[-1]

    assert lambda_max >= lambda_min
    spectral_norm = max(lambda_max, -lambda_min)

    evals, _ = torch.linalg.eigh(k)
    lambda_min = evals[0]
    lambda_max = evals[-1]

    assert lambda_max >= lambda_min
    ref_spectral_norm = max(lambda_max, -lambda_min)

    rel_spectral_norm = spectral_norm / ref_spectral_norm

    return spectral_norm, rel_spectral_norm
//  [python]
def spectral_clustering(Asym, k = 1):
    Dsqrt = Asym.sum(dim = -1).sqrt()
    normalized_symmetric_laplacian = torch.ones_like(Dsqrt).diag() - Asym / Dsqrt.unsqueeze(-1) / Dsqrt.unsqueeze(-2)
    eigval, eigvec = torch.symeig(normalized_symmetric_laplacian)
    return eigvec[:, k]
//  [python]
def spectral_clustering(Asym, k = 1):
    Dsqrt = Asym.sum(dim = -1).sqrt()
    normalized_symmetric_laplacian = torch.ones_like(Dsqrt).diag() - Asym / Dsqrt.unsqueeze(-1) / Dsqrt.unsqueeze(-2)
    eigval, eigvec = torch.linalg.eigh(normalized_symmetric_laplacian)
    return eigvec[:, k]
//  [python]
def spectral_order(Asym):
    unnormalized_symmetric_laplacian = Asym.sum(dim = -1).diag() - Asym
    Dsqrt = Asym.sum(dim = -1).sqrt()
    normalized_symmetric_laplacian = torch.ones_like(Dsqrt).diag() - Asym / Dsqrt.unsqueeze(-1) / Dsqrt.unsqueeze(-2)
    eigval, eigvec = torch.symeig(normalized_symmetric_laplacian)
    return eigvec[:, 1].argsort()
//  [python]
def spectral_order(Asym):
    unnormalized_symmetric_laplacian = Asym.sum(dim = -1).diag() - Asym
    Dsqrt = Asym.sum(dim = -1).sqrt()
    normalized_symmetric_laplacian = torch.ones_like(Dsqrt).diag() - Asym / Dsqrt.unsqueeze(-1) / Dsqrt.unsqueeze(-2)
    eigval, eigvec = torch.linalg.eigh(normalized_symmetric_laplacian)
    return eigvec[:, 1].argsort()
//  [python]
def get_gaussian_ellipsiod(
    mean: Float[Tensor, "3"],
    cov: Float[Tensor, "3 3"],
    n_std: int = 2,
    color="lightblue",
    opacity: float = 0.5,
    resolution: int = 20,
    name: str = "ellipse",
) -> go.Mesh3d:  # type: ignore
    """Get a plotly ellipsoid for a Gaussian.

    Args:
        mean: mean of the Gaussian.
        cov: covariance of the Gaussian.
        n_std: Standard deviation to visualize. Defaults to 2 (95% confidence).
        color: Color of the ellipsoid. Defaults to None.
        opacity: Opacity of the ellipsoid. Defaults to 0.5.
        resolution: Resolution of the ellipsoid. Defaults to 20.
        name: Name of the ellipsoid. Defaults to "ellipse".

    Returns:
        ellipsoid object.
    """

    phi = torch.linspace(0, 2 * torch.pi, resolution)
    theta = torch.linspace(-torch.pi / 2, torch.pi / 2, resolution)
    phi, theta = torch.meshgrid(phi, theta, indexing="ij")

    x = torch.cos(theta) * torch.sin(phi)
    y = torch.cos(theta) * torch.cos(phi)
    z = torch.sin(theta)
    pts = torch.stack((x, y, z), dim=-1)

    eigenvals, eigenvecs = torch.symeig(cov, eigenvectors=True)
    idx = torch.sum(cov, dim=0).argsort()
    idx = eigenvals[idx].argsort()
    eigenvals = eigenvals[idx]
    eigenvecs = eigenvecs[:, idx]

    scaling = torch.sqrt(eigenvals) * n_std
    pts = pts * scaling

    pts = pts @ eigenvecs.t()

    pts += mean

    return go.Mesh3d(  # type: ignore
        {
            "x": pts[:, :, 0].flatten(),
            "y": pts[:, :, 1].flatten(),
            "z": pts[:, :, 2].flatten(),
            "alphahull": 0,
            "opacity": opacity,
            "color": color,
            "name": name,
        }
    )
//  [python]
def get_gaussian_ellipsiod(
    mean: Float[Tensor, "3"],
    cov: Float[Tensor, "3 3"],
    n_std: int = 2,
    color="lightblue",
    opacity: float = 0.5,
    resolution: int = 20,
    name: str = "ellipse",
) -> go.Mesh3d:  # type: ignore
    """Get a plotly ellipsoid for a Gaussian.

    Args:
        mean: mean of the Gaussian.
        cov: covariance of the Gaussian.
        n_std: Standard deviation to visualize. Defaults to 2 (95% confidence).
        color: Color of the ellipsoid. Defaults to None.
        opacity: Opacity of the ellipsoid. Defaults to 0.5.
        resolution: Resolution of the ellipsoid. Defaults to 20.
        name: Name of the ellipsoid. Defaults to "ellipse".

    Returns:
        ellipsoid object.
    """

    phi = torch.linspace(0, 2 * torch.pi, resolution)
    theta = torch.linspace(-torch.pi / 2, torch.pi / 2, resolution)
    phi, theta = torch.meshgrid(phi, theta, indexing="ij")

    x = torch.cos(theta) * torch.sin(phi)
    y = torch.cos(theta) * torch.cos(phi)
    z = torch.sin(theta)
    pts = torch.stack((x, y, z), dim=-1)

    eigenvals, eigenvecs = torch.linalg.eigh(cov)
    idx = torch.sum(cov, dim=0).argsort()
    idx = eigenvals[idx].argsort()
    eigenvals = eigenvals[idx]
    eigenvecs = eigenvecs[:, idx]

    scaling = torch.sqrt(eigenvals) * n_std
    pts = pts * scaling

    pts = pts @ eigenvecs.t()

    pts += mean

    return go.Mesh3d(  # type: ignore
        {
            "x": pts[:, :, 0].flatten(),
            "y": pts[:, :, 1].flatten(),
            "z": pts[:, :, 2].flatten(),
            "alphahull": 0,
            "opacity": opacity,
            "color": color,
            "name": name,
        }
    )
//  [python]
def rot_to_quat(
    rot: Tensor,
):
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot

    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,
        ],
    ]
    print("about to einh")
    k = (1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)
    print(f"calling einh on {k.shape}")
    _, vectors = torch.symeig(k, eigenvectors=True)
    print("einh vecs",vectors.shape)
    return vectors[..., -1]
//  [python]
def rot_to_quat(
    rot: Tensor,
):
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot

    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,
        ],
    ]
    print("about to einh")
    k = (1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)
    print(f"calling einh on {k.shape}")
    _, vectors = torch.linalg.eigh(k)
    print("einh vecs",vectors.shape)
    return vectors[..., -1]
//  [python]
def compute_inv_covariance(centroids, train_features, jitters=None):
    r"""
    This function computes inverse covariance matrix that is required by Mahalanobis distance:
    MD = \sqrt((h(x) - \mu)^{T} \Sigma^{-1} (h(x) - \mu))

    """

    # jitter is the value to be added to the covariance matrix

    if jitters is None:
        jitters = JITTERS
    jitter = 0
    jitter_eps = None

    # A nested loop iterates over each centroid (mu_c) and the corresponding training features (x) for that centroid.
    # and for each pair of centroid and feature, the difference (d) between the feature and centroid is computed and
    # the outer product of d with itself is added to the covariance matrix.

    if torch.cuda.is_available():
        centroids = centroids.cuda()
        train_features = train_features.cuda()

    cov_scaled = torch.cov(train_features.T)

    # The function then iterates over each jitter_eps value in jitters and adds jitter to the scaled covariance matrix.
    # And the eigenvalues of the updated covariance matrix are computed, and if all eigenvalues are non-negative, the loop breaks.

    for i, jitter_eps in enumerate(jitters):
        jitter = jitter_eps * torch.eye(
            cov_scaled.shape[1],
            device=cov_scaled.device,
        )
        cov_scaled_update = cov_scaled + jitter
        eigenvalues = torch.symeig(cov_scaled_update).eigenvalues
        if (eigenvalues >= 0).all():
            break
    cov_scaled = cov_scaled + jitter

    # finally computes inverse of scaled covariance matrix with regularisation for MD calculation

    cov_inv = torch.inverse(cov_scaled.to(torch.float64)).float()
    return cov_inv, jitter_eps
//  [python]
def compute_inv_covariance(centroids, train_features, jitters=None):
    r"""
    This function computes inverse covariance matrix that is required by Mahalanobis distance:
    MD = \sqrt((h(x) - \mu)^{T} \Sigma^{-1} (h(x) - \mu))

    """

    # jitter is the value to be added to the covariance matrix

    if jitters is None:
        jitters = JITTERS
    jitter = 0
    jitter_eps = None

    # A nested loop iterates over each centroid (mu_c) and the corresponding training features (x) for that centroid.
    # and for each pair of centroid and feature, the difference (d) between the feature and centroid is computed and
    # the outer product of d with itself is added to the covariance matrix.

    if torch.cuda.is_available():
        centroids = centroids.cuda()
        train_features = train_features.cuda()

    cov_scaled = torch.cov(train_features.T)

    # The function then iterates over each jitter_eps value in jitters and adds jitter to the scaled covariance matrix.
    # And the eigenvalues of the updated covariance matrix are computed, and if all eigenvalues are non-negative, the loop breaks.

    for i, jitter_eps in enumerate(jitters):
        jitter = jitter_eps * torch.eye(
            cov_scaled.shape[1],
            device=cov_scaled.device,
        )
        cov_scaled_update = cov_scaled + jitter
        eigenvalues = torch.linalg.eigh(cov_scaled_update).eigenvalues
        if (eigenvalues >= 0).all():
            break
    cov_scaled = cov_scaled + jitter

    # finally computes inverse of scaled covariance matrix with regularisation for MD calculation

    cov_inv = torch.inverse(cov_scaled.to(torch.float64)).float()
    return cov_inv, jitter_eps
//  [python]
def matrix_decomposition(matrix=None, epsilon=1e-6, method='regularize'):

    """
    Eigen-decomposition for the input hermitian matrix.

    Parameters
    ----------
    matrix: torch.Tensor, default=None
        The hermitian matrix used to perform eigen-decomposition.
    epsilon: float, default=1e-6
        The numerical modification put on the matrix to ensure its positive rank and may used as cutoff for eigenvalues.
    method: string, default='regularize
        Use three different methods to do eigen-decomposition: 'regularize', 'cutoff', 'clamp'
        regularize: Modify the input matrix by adding epsilon driven identify matrix and take absolute value for eigenvalues;
        cutoff: Set epsilon as the cutoff for the eigenvalues of the matrix;
        clamp: Set all eigenvalues smaller than epsilon to be epsilon;
    
    Returns
    ----------
    eigenvalues, eigenvectors (torch.Tensor) of the input hermitian matrix.
    """

    methods = ['regularize', 'cutoff', 'clamp']
    assert method in methods, "Invalid method "+str(method)+", should use methods from "+str(methods)

    _matrix = matrix + torch.eye(matrix.shape[0]) * epsilon
    eigenval, eigenvec = torch.symeig(_matrix, eigenvectors=True)
    eigenvec = eigenvec.t()

    if method == "regularize":
        eigenval = torch.abs(eigenval)
#         eigenvec = torch.abs(eigenvec)
    elif method == "cutoff":
        _mask = eigenval > epsilon
        eigenval = eigenval[_mask]
        eigenvec = eigenvec[_mask]
    elif method == "clamp":
        eigval = torch.clamp_min(eigval, min=epsilon)
    return eigenval, eigenvec
//  [python]
def matrix_decomposition(matrix=None, epsilon=1e-6, method='regularize'):

    """
    Eigen-decomposition for the input hermitian matrix.

    Parameters
    ----------
    matrix: torch.Tensor, default=None
        The hermitian matrix used to perform eigen-decomposition.
    epsilon: float, default=1e-6
        The numerical modification put on the matrix to ensure its positive rank and may used as cutoff for eigenvalues.
    method: string, default='regularize
        Use three different methods to do eigen-decomposition: 'regularize', 'cutoff', 'clamp'
        regularize: Modify the input matrix by adding epsilon driven identify matrix and take absolute value for eigenvalues;
        cutoff: Set epsilon as the cutoff for the eigenvalues of the matrix;
        clamp: Set all eigenvalues smaller than epsilon to be epsilon;
    
    Returns
    ----------
    eigenvalues, eigenvectors (torch.Tensor) of the input hermitian matrix.
    """

    methods = ['regularize', 'cutoff', 'clamp']
    assert method in methods, "Invalid method "+str(method)+", should use methods from "+str(methods)

    _matrix = matrix + torch.eye(matrix.shape[0]) * epsilon
    eigenval, eigenvec = torch.linalg.eigh(_matrix)
    eigenvec = eigenvec.t()

    if method == "regularize":
        eigenval = torch.abs(eigenval)
#         eigenvec = torch.abs(eigenvec)
    elif method == "cutoff":
        _mask = eigenval > epsilon
        eigenval = eigenval[_mask]
        eigenvec = eigenvec[_mask]
    elif method == "clamp":
        eigval = torch.clamp_min(eigval, min=epsilon)
    return eigenval, eigenvec
//  [python]
def forward(self, x):  # x: [B,C,T]
        # x is not NORMALIZED

        if self.logcov:
            # compute log-covariance
            # log(C) = Ulog(V)U.T
            xm = x - x.mean(dim=2, keepdim=True)
            cov = torch.matmul(xm, xm.transpose(2, 1)) / (xm.size(2) - 1)

            e, v = torch.symeig(cov, eigenvectors=True)
            e = torch.log(torch.clamp(e, min=1e-10))
            log_cov = torch.matmul(
                torch.matmul(v, e.diag_embed()), v.transpose(2, 1)
            )  # [B,C,C]

            feats = log_cov[:, self.inds[0], self.inds[1]]  # [B,n_in*(n_in+1)/2]

        else:
            # compute log-variance
            feats = torch.log(torch.var(x, dim=2, unbiased=True))
            feats[torch.isneginf(feats)] = 0

        feats = self.linear(feats)  # [B,n_out*n_in+n_out]

        w = feats[:, : -self.n_out_channels].reshape(
            -1, self.n_out_channels, self.n_in_channels
        )  # [B,n_out,n_in]
        b = feats[:, -self.n_out_channels :].reshape(
            -1, self.n_out_channels, 1
        )  # [B,n_out,1]

        if self.logcov and self.soft_thresh > 0:
            # apply soft-thresholding in [-thr,thr]
            w = torch.clamp(w - self.soft_thresh, min=0) + torch.clamp(
                w + self.soft_thresh, max=0
            )

        x = (x - x.mean(dim=2, keepdim=True)) / x.std(dim=2, keepdim=True)
        x = torch.matmul(w, x) + b

        return x
//  [python]
def forward(self, x):  # x: [B,C,T]
        # x is not NORMALIZED

        if self.logcov:
            # compute log-covariance
            # log(C) = Ulog(V)U.T
            xm = x - x.mean(dim=2, keepdim=True)
            cov = torch.matmul(xm, xm.transpose(2, 1)) / (xm.size(2) - 1)

            e, v = torch.linalg.eigh(cov, UPLO="U")
            e = torch.log(torch.clamp(e, min=1e-10))
            log_cov = torch.matmul(
                torch.matmul(v, e.diag_embed()), v.transpose(2, 1)
            )  # [B,C,C]

            feats = log_cov[:, self.inds[0], self.inds[1]]  # [B,n_in*(n_in+1)/2]

        else:
            # compute log-variance
            feats = torch.log(torch.var(x, dim=2, unbiased=True))
            feats[torch.isneginf(feats)] = 0

        feats = self.linear(feats)  # [B,n_out*n_in+n_out]

        w = feats[:, : -self.n_out_channels].reshape(
            -1, self.n_out_channels, self.n_in_channels
        )  # [B,n_out,n_in]
        b = feats[:, -self.n_out_channels :].reshape(
            -1, self.n_out_channels, 1
        )  # [B,n_out,1]

        if self.logcov and self.soft_thresh > 0:
            # apply soft-thresholding in [-thr,thr]
            w = torch.clamp(w - self.soft_thresh, min=0) + torch.clamp(
                w + self.soft_thresh, max=0
            )

        x = (x - x.mean(dim=2, keepdim=True)) / x.std(dim=2, keepdim=True)
        x = torch.matmul(w, x) + b

        return x
//  [python]
def U(self, lap_type: str = "sym"):
        lap = self.L(lap_type)
        if self._U is not None:
            U = self._U
        else:
            fs, U = torch.eig(lap.to_dense())
            if self.cache:
                fs[fs.abs() < 1e-6] = 0
                self._fs = fs
                self._U = U
        return U
//  [python]
def U(self, lap_type: str = "sym"):
        lap = self.L(lap_type)
        if self._U is not None:
            U = self._U
        else:
            fs, U = torch.linalg.eigh(lap.to_dense())
            if self.cache:
                fs[fs.abs() < 1e-6] = 0
                self._fs = fs
                self._U = U
        return U