code: def _prob_QK(self, Q, K, sample_k, n_top):  # n_top: c*ln(L_q)
      # Q [B, H, L, D]
      B, H, L_K, E = K.shape
      _, _, L_Q, _ = Q.shape
  
      # calculate the sampled Q_K
      K_expand = K.unsqueeze(-3).expand(B, H, L_Q, L_K, E)
      index_sample = torch.randint(
          L_K, (L_Q, sample_k)
      )  # real U = U_part(factor*ln(L_k))*L_q
      K_sample = K_expand[
          :, :, torch.arange(L_Q).unsqueeze(1), index_sample, :
      ]
      Q_K_sample = torch.matmul(
          Q.unsqueeze(-2), K_sample.transpose(-2, -1)
      ).squeeze()
  
      # find the Top_k query with sparisty measurement
      M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L_K)
      M_top = M.topk(n_top, sorted=False)[1]
  
      # use the reduced Q to calculate Q_K
      Q_reduce = Q[
          torch.arange(B)[:, None, None],
          torch.arange(H)[None, :, None],
          M_top,
          :,
      ]  # factor*ln(L_q)
      Q_K = torch.matmul(
          Q_reduce, K.transpose(-2, -1)
      )  # factor*ln(L_q)*L_k
  
      return Q_K, M_top
  
imports: import os
  import math
  import numpy
  from math import sqrt
  from scipy.fftpack import next_fast_len
label_desc: Use image node_id 2 for calling the ProbAttention obj's underlying member method code with example usage: obj._prob_QK(Q, K, sample_k, n_top) and returns: Q_K,  M_top