code: def convert(logits):
      "move root arcs from diagonal"
      new_shape = list(logits.shape)
      new_shape[1] += 1
      new_shape[2] += 1
      new_logits = torch.zeros(new_shape).type_as(logits.data)
      new_logits.fill_(-1e9)
      new_logits[:, 1:, 1:] = logits
  
      N = logits.size(1)
      new_logits[:, 0, 1:] = logits[:, torch.arange(N), torch.arange(N)]
      new_logits[
          :, torch.arange(1, N + 1), torch.arange(1, N + 1)
      ] = -1e9
      return new_logits
  
file_path: pytorch-struct/torch_struct/deptree.py
imports: import torch
  from .helpers import _Struct, Chart
label_desc: Use image node_id 1 for calling a global function with example usage: convert(logits) and returns: new_logits