text: import torch
  import numpy as np
  from torch_sparse import SparseTensor
  from thgsp.convert import to_xcipy, from_cpx, get_array_module, get_ddd, spmatrix
  
  
  def normalize_laplace(L: SparseTensor, lam_max: float = 2.):
      Ln = L.clone()
      row, col, val = Ln.coo()
      diag_mask = row == col
      val[...] = (2. * val) / lam_max
      val.masked_fill_(val == float('inf'), 0)
      val[diag_mask] -= 1
      return Ln
  
  
  def cheby_op(x: torch.Tensor, L: SparseTensor, coeff: torch.Tensor, lam_max: float = 2.):
      r""" Chebyshev approximation of graph filtering
  
      Parameters
      ----------
      x:          Tensor
          The input graph signal. It's shape can be either :obj:`(N,)` , :obj:`(N,Ci)` or :obj:`(Co,N,Ci)`, wherein
          :obj:`N`, :obj:`Ci` and :obj:`Co` are the numbers of nodes, input channels, and output channels respectively.
      L:          SparseTensor
          The :obj:`(N,N)` Laplacian matrix.
      coeff:      Tensor
          The :obj:`(Co,Ci,K+1)` Chebyshev coefficients for :obj:`Ci*Co` kernels, wherein :obj:`K` is the order of
          approximation.
      lam_max:    float,optional
          The maximal graph frequency, i.e., :math:`\lambda_{max}`
  
      Returns
      -------
      Tensor
          The filtered signals of shape :obj:`(Co,N,Ci)`
      """
  
      Co, Ci, K = coeff.shape
      N = L.size(-1)
  
      if x.dim() == 1:
          assert x.size() == (N,)
          x = x[None, ..., None]  # (N,) --> 1 x N x 1
      elif x.dim() == 2:  # N x Ci --> Co x N x Ci
          assert x.size() == (N, Ci)
          x = x.unsqueeze(0)
      elif x.dim() == 3:  # Co x N x Ci
          assert x.size() == (Co, N, Ci) or (1, N, Ci)
      else:
          raise RuntimeError(
              "The input signals has mismatched dimensions: {}".format(x.size()))
  
      K = K - 1
      c = coeff.unsqueeze(1)  # Co x Ci x K --> Co x 1 x Ci x K
      L_norm = normalize_laplace(L, lam_max)
      twf_old = x
      twf_cur = L_norm @ x  # Co x N x Ci
      result = 0.5 * c[..., 0] * twf_old + c[..., 1] * twf_cur
      for k in range(2, K + 1):
          twf_new = 2 * (L_norm @ twf_cur) - twf_old
          result = result + c[..., k] * twf_new
          twf_old = twf_cur
          twf_cur = twf_new
  
      return result
  
  
  def cheby_op_basis(L, coeff, lam_max=2., return_ts=False):
      assert coeff.ndim == 1
      K = len(coeff)
      if isinstance(L, SparseTensor):
          dt, dv, density, on_gpu = get_ddd(L)
          xp, xcipy, _ = get_array_module(on_gpu)
      elif isinstance(L, spmatrix):
          import scipy
          xcipy = scipy
          xp = np
      else:
          raise TypeError
  
      L = to_xcipy(L)
      N = L.shape[-1]
  
      coeff = xp.asarray(coeff)
      I = xcipy.sparse.eye(N, dtype=L.dtype, format="csr")
  
      Ln = L * (2 / lam_max) - I
      Tl_old = I
      Tl_cur = Ln
      Hl = 0.5 * coeff[0] * Tl_old + coeff[1] * Tl_cur
      for k in range(2, K):
          Tl_new = 2 * Ln * Tl_cur - Tl_old
          Hl = Hl + coeff[k] * Tl_new
          Tl_old = Tl_cur
          Tl_cur = Tl_new
  
      if return_ts:
          result = SparseTensor.from_scipy(Hl) if xp == np else from_cpx(Hl)
      else:
          result = Hl
      return result
  
  
  def cheby_coeff(kernels, K=10, lam_max=2., num_points=None, dtype=None, device=None):
      if num_points is None:
          num_points = K + 1
      assert lam_max > 0
  
      if isinstance(kernels, np.ndarray):
          M, Co, Ci = kernels.shape
      else:  # pass only a single kernel function
          kernels = np.array([[[kernels]]])  # 1 x 1 x 1 array
          M = Co = Ci = 1
  
      points = np.pi * (torch.arange(num_points, dtype=dtype,
                                     device=device) + 0.5) / num_points
  
      gs = torch.zeros(M, Co, Ci, 1, num_points, dtype=dtype, device=device)
      kernel_cache = dict()
      for m in range(M):
          for i in range(Ci):
              for j in range(Co):
                  krn = kernels[m, j, i]
                  kid = id(krn)
                  if kid in kernel_cache:
                      gs[m, j, i] = kernel_cache[kid]
                  else:  # kernel_cache stores the reference to a part of memory used by gs.
                      gs[m, j, i] = krn(lam_max / 2 * (torch.cos(points) + 1))
                      kernel_cache[kid] = gs[m, j, i]
      order_cos = torch.arange(
          K + 1, dtype=dtype, device=device).reshape(1, -1) * points.reshape(-1, 1)
      order_cos.cos_()  # num x K+1
      coeff = (gs @ order_cos) * 2. / num_points  # M x Co x Ci x 1 x K+1
      return coeff.squeeze_(-2)  # M x Co x Ci x K+1
  
  
  def polyval(c, x):
      """
      Evaluate `N`-order polynomial at the points `x` with the given `N+1` coefficients
      in descending order.
  
      Parameters
      ----------
      c:  Tensor, a list of Tensor
          The coefficients can be either a list of  tensors which are broadcastable with the input `x`
          or an concatenated tensor of them.
  
      x:  Tensor, scalar
          Arbitrary rank-D tensor(D=1,2,...) is valid due to the broadcasting semantics.
  
      Returns
      -------
      Tensor
             A tensor with the same shape as `x`
      """
      order = len(c) - 1
      f = c[0] * x + c[1]
      for i in range(2, order + 1):
          f = f * x + c[i]
      return f
  
  
  def nla(x, frac=0.4, k=None, scheme='abs'):
      Co, N, Ci = x.shape
  
      if k is not None:
          k_largest = k
      else:
          k_largest = int(frac * N)
  
      fuse = x.reshape(-1, Ci)
      if scheme == 'abs':
          _, idx = fuse.abs().topk(k_largest, dim=0)
          val = fuse.gather(0, idx)
      elif scheme == 'naive':
          val, idx = fuse.topk(k_largest, dim=0)
  
      elif scheme == "keeplow":
          fuse_high = fuse[N:, ...]
          _, idx = fuse_high.abs().topk(k_largest, dim=0)
          val = fuse_high.gather(0, idx)
          res = fuse_high.new_zeros((Co - 1) * N, Ci)
          res.scatter_(0, idx, val)
          return torch.cat([x[0, ...], res]).reshape(Co, N, Ci)
  
      else:
          raise RuntimeError(
              "{} is not a valid supported non-linear approximation scheme".format(scheme))
      res = x.new_zeros(Co * N, Ci)
      res.scatter_(0, idx, val)
      return res.reshape(Co, N, Ci)
  
  
  def hard_threshold(x, T=0.3, lowest=False):
      if not lowest:
          x[x.abs() < T] = 0
      else:
          x_high_pass = x[1:, ...]
          x_high_pass[x_high_pass.abs() < T] = 0
      return x
  