docstr: Compute the marginals of a CFG using CKY.
  
  Parameters:
      scores : terms : b x n x T
               rules : b x NT x (NT+T) x (NT+T)
               root:   b x NT
      lengths : lengths in batch
  
  Returns:
      v: b tensor of total sum
      spans: bxNxT terms, (bxNTx(NT+S)x(NT+S)) rules, bxNT roots
docstr_tok: ["Compute","the","marginals","of","a","CFG","using","CKY",".","Parameters",":","scores",":","terms",":","b","x","n","x","T","rules",":","b","x","NT","x","(","NT+T",")","x","(","NT+T",")","root",":","b","x","NT","lengths",":","lengths","in","batch","Returns",":","v",":","b","tensor","of","total","sum","spans",":","bxNxT","terms",",","(","bxNTx","(","NT+S",")","x","(","NT+S",")",")","rules",",","bxNT","roots"]
code: def marginals(self, scores, lengths=None, _autograd=True, _raw=False):
      """
      Compute the marginals of a CFG using CKY.
  
      Parameters:
          scores : terms : b x n x T
                   rules : b x NT x (NT+T) x (NT+T)
                   root:   b x NT
          lengths : lengths in batch
  
      Returns:
          v: b tensor of total sum
          spans: bxNxT terms, (bxNTx(NT+S)x(NT+S)) rules, bxNT roots
  
      """
      terms, rules, roots = scores
      batch, N, T = terms.shape
      _, NT, _, _ = rules.shape
  
      v, (term_use, rule_use, root_use, spans) = self.logpartition(
          scores, lengths=lengths, force_grad=True
      )
  
      def marginal(obj, inputs):
          obj = self.semiring.unconvert(obj).sum(dim=0)
          marg = torch.autograd.grad(
              obj,
              inputs,
              create_graph=True,
              only_inputs=True,
              allow_unused=False,
          )
  
          spans_marg = torch.zeros(
              batch,
              N,
              N,
              NT,
              dtype=scores[1].dtype,
              device=scores[1].device,
          )
          span_ls = marg[3:]
          for w in range(len(span_ls)):
              x = span_ls[w].sum(dim=0, keepdim=True)
              spans_marg[:, w, : N - w - 1] = self.semiring.unconvert(x)
  
          rule_marg = self.semiring.unconvert(marg[0]).squeeze(1)
          root_marg = self.semiring.unconvert(marg[1])
          term_marg = self.semiring.unconvert(marg[2])
  
          assert term_marg.shape == (batch, N, T)
          assert root_marg.shape == (batch, NT)
          assert rule_marg.shape == (batch, NT, NT + T, NT + T)
          return (term_marg, rule_marg, root_marg, spans_marg)
  
      inputs = (rule_use, root_use, term_use) + tuple(spans)
      if _raw:
          paths = []
          for k in range(v.shape[0]):
              obj = v[k : k + 1]
              marg = marginal(obj, inputs)
              paths.append(marg[-1])
          paths = torch.stack(paths, 0)
          obj = v.sum(dim=0, keepdim=True)
          term_marg, rule_marg, root_marg, _ = marginal(obj, inputs)
          return term_marg, rule_marg, root_marg, paths
      else:
          return marginal(v, inputs)
  
file_path: pytorch-struct/torch_struct/cky.py
imports: import torch
  from .helpers import _Struct, Chart
label_desc: Use image node_id 2 for calling the CKY obj's underlying member method code with example usage: obj.marginals(scores, lengths, _autograd, _raw) and returns: term_marg,  rule_marg,  root_marg,  spans_marg, term_marg,  rule_marg,  root_marg,  paths, marginal