code: def logpartition(self, scores, lengths=None, force_grad=False):
      semiring = self.semiring
  
      # Checks
      terms, rules, roots = scores
      rules.requires_grad_(True)
      ssize = semiring.size()
      batch, N, T = terms.shape
      _, NT, _, _ = rules.shape
      S = NT + T
  
      terms, rules, roots = (
          semiring.convert(terms).requires_grad_(True),
          semiring.convert(rules).requires_grad_(True),
          semiring.convert(roots).requires_grad_(True),
      )
      if lengths is None:
          lengths = torch.LongTensor([N] * batch).to(terms.device)
  
      # Charts
      beta = [
          Chart((batch, N, N, NT), rules, semiring) for _ in range(2)
      ]
      span = [None for _ in range(N)]
      v = (ssize, batch)
      term_use = terms + 0.0
  
      # Split into NT/T groups
      NTs = slice(0, NT)
      Ts = slice(NT, S)
      rules = rules.view(ssize, batch, 1, NT, S, S)
  
      def arr(a, b):
          return (
              rules[..., a, b]
              .contiguous()
              .view(*v + (NT, -1))
              .transpose(-2, -1)
          )
  
      matmul = semiring.matmul
      times = semiring.times
      X_Y_Z = arr(NTs, NTs)
      X_Y1_Z = arr(Ts, NTs)
      X_Y_Z1 = arr(NTs, Ts)
      X_Y1_Z1 = arr(Ts, Ts)
  
      for w in range(1, N):
          all_span = []
          v2 = v + (N - w, -1)
  
          Y = beta[A][: N - w, :w, :]
          Z = beta[B][w:, N - w :, :]
          X1 = matmul(matmul(Y.transpose(-2, -1), Z).view(*v2), X_Y_Z)
          all_span.append(X1)
  
          Y_term = term_use[..., : N - w, :, None]
          Z_term = term_use[..., w:, None, :]
  
          Y = Y[..., -1, :].unsqueeze(-1)
          X2 = matmul(times(Y, Z_term).view(*v2), X_Y_Z1)
  
          Z = Z[..., 0, :].unsqueeze(-2)
          X3 = matmul(times(Y_term, Z).view(*v2), X_Y1_Z)
          all_span += [X2, X3]
  
          if w == 1:
              X4 = matmul(times(Y_term, Z_term).view(*v2), X_Y1_Z1)
              all_span.append(X4)
  
          span[w] = semiring.sum(torch.stack(all_span, dim=-1))
          beta[A][: N - w, w, :] = span[w]
          beta[B][w:N, N - w - 1, :] = span[w]
  
      final = beta[A][0, :, NTs]
      top = torch.stack(
          [final[:, i, l - 1] for i, l in enumerate(lengths)], dim=1
      )
      log_Z = semiring.dot(top, roots)
      return log_Z, (term_use, rules, roots, span[1:])
  
file_path: pytorch-struct/torch_struct/cky.py
imports: import torch
  from .helpers import _Struct, Chart
label_desc: Use image node_id 1 for calling the CKY obj's underlying member method code with example usage: obj.logpartition(scores, lengths, force_grad) and returns: log_Z,  term_use,  rules,  roots,  span, rules