code: def deptree_part(arc_scores, multi_root, lengths=None, eps=1e-5):
      if lengths is not None:
          batch, N, N = arc_scores.shape
          x = torch.arange(N, device=arc_scores.device).expand(batch, N)
          if not torch.is_tensor(lengths):
              lengths = torch.tensor(lengths, device=arc_scores.device)
          lengths = lengths.unsqueeze(1)
          x = x < lengths
          det_offset = torch.diag_embed((~x).float())
          x = x.unsqueeze(2).expand(-1, -1, N)
          mask = torch.transpose(x, 1, 2) * x
          mask = mask.float()
          mask[mask == 0] = float("-inf")
          mask[mask == 1] = 0
          arc_scores = arc_scores + mask
      input = arc_scores
      eye = torch.eye(input.shape[1], device=input.device)
      laplacian = input.exp() + eps
      lap = laplacian.masked_fill(eye != 0, 0)
      lap = -lap + torch.diag_embed(
          lap.sum(1), offset=0, dim1=-2, dim2=-1
      )
      if lengths is not None:
          lap += det_offset
  
      if multi_root:
          rss = torch.diagonal(
              input, 0, -2, -1
          ).exp()  # root selection scores
          lap = lap + torch.diag_embed(rss, offset=0, dim1=-2, dim2=-1)
      else:
          lap[:, 0] = torch.diagonal(input, 0, -2, -1).exp()
      return lap.logdet()
  
file_path: pytorch-struct/torch_struct/deptree.py
imports: import torch
  from .helpers import _Struct, Chart
label_desc: Use image node_id 3 for calling a global function with example usage: deptree_part(arc_scores, multi_root, lengths, eps) and returns: lap