code: def extract_topk(matrix, lengths):
      batch, K, N = matrix.shape[:3]
      spans = []
      trees = []
      for b in range(batch):
          for k in range(K):
              this_span = matrix[b][k]
              span, tree = extract_parse(this_span, lengths[b])
              trees.append(tree)
              spans.append(span)
              # print(span)
              # print(tree)
          # break
      return spans, trees
  
imports: import torch
  from torch_struct import SentCFG
label_desc: Use image node_id 3 for calling a global function with example usage: extract_topk(matrix, lengths) and returns: spans,  trees