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