code: def expand_spans(spans, words, K, V):
      batch, N = words.shape
      spans[:, torch.arange(N), torch.arange(N)] = 0
      new_spans = torch.zeros(K, batch, N, N, 1 + V).cuda()
      new_spans[:, :, :, :, :1] = spans.view(K, batch, N, N, 1)
      new_spans[:, :, torch.arange(N), torch.arange(N), :].fill_(0)
      new_spans[:, :, torch.arange(N), torch.arange(N), 1:] = (
          torch.nn.functional.one_hot(words, V)
          .float()
          .cuda()
          .view(1, batch, N, V)
      )
      new_spans = new_spans.view(batch * K, N, N, 1 + V)
      return new_spans
  
file_path: pytorch-struct/examples/tree.py
imports: from torch_struct import TreeCRF, SelfCritical
  import torchtext.data
  from torch_struct.data import ListOpsDataset, TokenBucket
  from torch_struct.networks import TreeLSTM, SpanLSTM
  import torch
  import torch.nn
  import wandb
label_desc: Use image node_id 2 for calling a global function with example usage: expand_spans(spans, words, K, V) and returns: new_spans