class_docstr: Create a graph pool in advance to accelerate graph building phase in Transformer.
docstr: Return a batched graph for beam search during inference of Transformer.
  args:
      src_buf: a list of input sequence
      start_sym: the index of start-of-sequence symbol
      max_len: maximum length for decoding
      k: beam size
      device: 'cpu' or 'cuda:*'
docstr_tok: ["Return","a","batched","graph","for","beam","search","during","inference","of","Transformer",".","args",":","src_buf",":","a","list","of","input","sequence","start_sym",":","the","index","of","start-of-sequence","symbol","max_len",":","maximum","length","for","decoding","k",":","beam","size","device",":","'cpu","'","or","'cuda",":","*","'"]
code: def beam(self, src_buf, start_sym, max_len, k, device="cpu"):
      """
      Return a batched graph for beam search during inference of Transformer.
      args:
          src_buf: a list of input sequence
          start_sym: the index of start-of-sequence symbol
          max_len: maximum length for decoding
          k: beam size
          device: 'cpu' or 'cuda:*'
      """
      g_list = []
      src_lens = [len(_) for _ in src_buf]
      tgt_lens = [max_len] * len(src_buf)
      num_edges = {"ee": [], "ed": [], "dd": []}
      for src_len, tgt_len in zip(src_lens, tgt_lens):
          i, j = src_len - 1, tgt_len - 1
          for _ in range(k):
              g_list.append(self.g_pool[i][j])
          for key in ["ee", "ed", "dd"]:
              num_edges[key].append(int(self.num_edges[key][i][j]))
  
      g = dgl.batch(g_list)
      src, tgt = [], []
      src_pos, tgt_pos = [], []
      enc_ids, dec_ids = [], []
      e2e_eids, e2d_eids, d2d_eids = [], [], []
      n_nodes, n_edges, n_tokens = 0, 0, 0
      for src_sample, n, n_ee, n_ed, n_dd in zip(
          src_buf,
          src_lens,
          num_edges["ee"],
          num_edges["ed"],
          num_edges["dd"],
      ):
          for _ in range(k):
              src.append(
                  th.tensor(src_sample, dtype=th.long, device=device)
              )
              src_pos.append(th.arange(n, dtype=th.long, device=device))
              enc_ids.append(
                  th.arange(
                      n_nodes, n_nodes + n, dtype=th.long, device=device
                  )
              )
              n_nodes += n
              e2e_eids.append(
                  th.arange(
                      n_edges,
                      n_edges + n_ee,
                      dtype=th.long,
                      device=device,
                  )
              )
              n_edges += n_ee
              tgt_seq = th.zeros(max_len, dtype=th.long, device=device)
              tgt_seq[0] = start_sym
              tgt.append(tgt_seq)
              tgt_pos.append(
                  th.arange(max_len, dtype=th.long, device=device)
              )
  
              dec_ids.append(
                  th.arange(
                      n_nodes,
                      n_nodes + max_len,
                      dtype=th.long,
                      device=device,
                  )
              )
              n_nodes += max_len
              e2d_eids.append(
                  th.arange(
                      n_edges,
                      n_edges + n_ed,
                      dtype=th.long,
                      device=device,
                  )
              )
              n_edges += n_ed
              d2d_eids.append(
                  th.arange(
                      n_edges,
                      n_edges + n_dd,
                      dtype=th.long,
                      device=device,
                  )
              )
              n_edges += n_dd
  
      g.set_n_initializer(dgl.init.zero_initializer)
      g.set_e_initializer(dgl.init.zero_initializer)
      g = g.to(device).long()
  
      return Graph(
          g=g,
          src=(th.cat(src), th.cat(src_pos)),
          tgt=(th.cat(tgt), th.cat(tgt_pos)),
          tgt_y=None,
          nids={"enc": th.cat(enc_ids), "dec": th.cat(dec_ids)},
          eids={
              "ee": th.cat(e2e_eids),
              "ed": th.cat(e2d_eids),
              "dd": th.cat(d2d_eids),
          },
          nid_arr={"enc": enc_ids, "dec": dec_ids},
          n_nodes=n_nodes,
          n_edges=n_edges,
          n_tokens=n_tokens,
      )
  
file_path: dgl/examples/pytorch/transformer/dataset/graph.py
imports: import itertools
  import time
  from collections import *
  import numpy
  import torch
  import dgl
label_desc: Use image node_id 2 for calling the GraphPool obj's underlying member method code with example usage: obj.beam(src_buf, start_sym, max_len, k, device) and returns: Graph