code: def augment_edge(data):
      """
      Input:
          data: PyG data object
      Output:
          data (edges are augmented in the following ways):
              data.edge_index: Added next-token edge. The inverse edges were also added.
              data.edge_attr (torch.Long):
                  data.edge_attr[:,0]: whether it is AST edge (0) for next-token edge (1)
                  data.edge_attr[:,1]: whether it is original direction (0) or inverse direction (1)
      """
  
      ##### AST edge
      edge_index_ast = data.edge_index
      edge_attr_ast = torch.zeros((edge_index_ast.size(1), 2))
  
      ##### Inverse AST edge
      edge_index_ast_inverse = torch.stack(
          [edge_index_ast[1], edge_index_ast[0]], dim=0
      )
      edge_attr_ast_inverse = torch.cat(
          [
              torch.zeros(edge_index_ast_inverse.size(1), 1),
              torch.ones(edge_index_ast_inverse.size(1), 1),
          ],
          dim=1,
      )
  
      ##### Next-token edge
  
      ## Obtain attributed nodes and get their indices in dfs order
      # attributed_node_idx = torch.where(data.node_is_attributed.view(-1,) == 1)[0]
      # attributed_node_idx_in_dfs_order = attributed_node_idx[torch.argsort(data.node_dfs_order[attributed_node_idx].view(-1,))]
  
      ## Since the nodes are already sorted in dfs ordering in our case, we can just do the following.
      attributed_node_idx_in_dfs_order = torch.where(
          data.node_is_attributed.view(
              -1,
          )
          == 1
      )[0]
  
      ## build next token edge
      # Given: attributed_node_idx_in_dfs_order
      #        [1, 3, 4, 5, 8, 9, 12]
      # Output:
      #    [[1, 3, 4, 5, 8, 9]
      #     [3, 4, 5, 8, 9, 12]
      edge_index_nextoken = torch.stack(
          [
              attributed_node_idx_in_dfs_order[:-1],
              attributed_node_idx_in_dfs_order[1:],
          ],
          dim=0,
      )
      edge_attr_nextoken = torch.cat(
          [
              torch.ones(edge_index_nextoken.size(1), 1),
              torch.zeros(edge_index_nextoken.size(1), 1),
          ],
          dim=1,
      )
  
      ##### Inverse next-token edge
      edge_index_nextoken_inverse = torch.stack(
          [edge_index_nextoken[1], edge_index_nextoken[0]], dim=0
      )
      edge_attr_nextoken_inverse = torch.ones(
          (edge_index_nextoken.size(1), 2)
      )
  
      data.edge_index = torch.cat(
          [
              edge_index_ast,
              edge_index_ast_inverse,
              edge_index_nextoken,
              edge_index_nextoken_inverse,
          ],
          dim=1,
      )
      data.edge_attr = torch.cat(
          [
              edge_attr_ast,
              edge_attr_ast_inverse,
              edge_attr_nextoken,
              edge_attr_nextoken_inverse,
          ],
          dim=0,
      )
  
      return data
  
file_path: ogb/examples/graphproppred/code2/utils.py
imports: import torch
  from collections import Counter
  import numpy
  import torch
label_desc: Use image node_id 2 for calling a global function with example usage: augment_edge(data) and returns: data