def get_k_q(q_k_mask):
    """Get the key-query index from query-key index for PAM-TVM"""
    k_q_mask = q_k_mask.clone()
    for i in range(len(q_k_mask)):
        for j in range(len(q_k_mask[0])):
            if q_k_mask[i, j] >= 0:
                k_q_mask[i, j] = torch.where(q_k_mask[q_k_mask[i, j]] ==i )[0]
    
    return k_q_mask
Get the key-query index from query-key index for PAM-TVM
PatchTST_supervised/Formers/Pyraformer/pyraformer/graph_attention.py