code: def to_parts(spans, extra, lengths=None):
      NT, T = extra
  
      batch, N, N, S = spans.shape
      assert S == NT + T
      terms = torch.zeros(batch, N, T)
      rules = torch.zeros(batch, NT, S, S)
      roots = torch.zeros(batch, NT)
      for b in range(batch):
          roots[b, :] = spans[b, 0, lengths[b] - 1, :NT]
          terms[b, : lengths[b]] = spans[
              b, torch.arange(lengths[b]), torch.arange(lengths[b]), NT:
          ]
          cover = spans[b].nonzero()
          left = {i: [] for i in range(N)}
          right = {i: [] for i in range(N)}
          for i in range(cover.shape[0]):
              i, j, A = cover[i].tolist()
              left[i].append((A, j, j - i + 1))
              right[j].append((A, i, j - i + 1))
          for i in range(cover.shape[0]):
              i, j, A = cover[i].tolist()
              B = None
              for B_p, k, a_span in left[i]:
                  for C_p, k_2, b_span in right[j]:
                      if k_2 == k + 1 and a_span + b_span == j - i + 1:
                          B, C = B_p, C_p
                          break
              if j > i:
                  assert B is not None, "%s" % (
                      (i, j, left[i], right[j], cover),
                  )
                  rules[b, A, B, C] += 1
      return terms, rules, roots
  
file_path: pytorch-struct/torch_struct/cky.py
imports: import torch
  from .helpers import _Struct, Chart
label_desc: Use image node_id 4 for calling the CKY obj's underlying member method code with example usage: obj.to_parts(spans, extra, lengths) and returns: terms,  rules,  roots