docstr: This function implements Beam Search in DGL, which is required in inference phase.
  Length normalization is given by (5 + len) ^ alpha / 6 ^ alpha. Please refer to https://arxiv.org/pdf/1609.08144.pdf.
  args:
      graph: a `Graph` object defined in `dgl.contrib.transformer.graph`.
      max_len: the maximum length of decoding.
      eos_id: the index of end-of-sequence symbol.
      k: beam size
  return:
      ret: a list of index array correspond to the input sequence specified by `graph``.
docstr_tok: ["This","function","implements","Beam","Search","in","DGL",",","which","is","required","in","inference","phase",".","Length","normalization","is","given","by","(","5","+","len",")","^","alpha","\/","6","^","alpha",".","Please","refer","to","https",":","\/\/arxiv.org\/pdf\/1609.08144.pdf",".","args",":","graph",":","a","`","Graph","`","object","defined","in","`","dgl.contrib.transformer.graph","`",".","max_len",":","the","maximum","length","of","decoding",".","eos_id",":","the","index","of","end-of-sequence","symbol",".","k",":","beam","size","return",":","ret",":","a","list","of","index","array","correspond","to","the","input","sequence","specified","by","`","graph","``","."]
code: def infer(self, graph, max_len, eos_id, k, alpha=1.0):
      """
      This function implements Beam Search in DGL, which is required in inference phase.
      Length normalization is given by (5 + len) ^ alpha / 6 ^ alpha. Please refer to https://arxiv.org/pdf/1609.08144.pdf.
      args:
          graph: a `Graph` object defined in `dgl.contrib.transformer.graph`.
          max_len: the maximum length of decoding.
          eos_id: the index of end-of-sequence symbol.
          k: beam size
      return:
          ret: a list of index array correspond to the input sequence specified by `graph``.
      """
      g = graph.g
      N, E = graph.n_nodes, graph.n_edges
      nids, eids = graph.nids, graph.eids
  
      # embed & pos
      src_embed = self.src_embed(graph.src[0])
      src_pos = self.pos_enc(graph.src[1])
      g.nodes[nids["enc"]].data["pos"] = graph.src[1]
      g.nodes[nids["enc"]].data["x"] = self.pos_enc.dropout(
          src_embed + src_pos
      )
      tgt_pos = self.pos_enc(graph.tgt[1])
      g.nodes[nids["dec"]].data["pos"] = graph.tgt[1]
  
      # init mask
      device = next(self.parameters()).device
      g.ndata["mask"] = th.zeros(N, dtype=th.uint8, device=device)
  
      # encode
      for i in range(self.encoder.N):
          pre_func = self.encoder.pre_func(i, "qkv")
          post_func = self.encoder.post_func(i)
          nodes, edges = nids["enc"], eids["ee"]
          self.update_graph(
              g, edges, [(pre_func, nodes)], [(post_func, nodes)]
          )
  
      # decode
      log_prob = None
      y = graph.tgt[0]
      for step in range(1, max_len):
          y = y.view(-1)
          tgt_embed = self.tgt_embed(y)
          g.ndata["x"][nids["dec"]] = self.pos_enc.dropout(
              tgt_embed + tgt_pos
          )
          edges_ed = g.filter_edges(
              lambda e: (e.dst["pos"] < step) & ~e.dst["mask"].bool(),
              eids["ed"],
          )
          edges_dd = g.filter_edges(
              lambda e: (e.dst["pos"] < step) & ~e.dst["mask"].bool(),
              eids["dd"],
          )
          nodes_d = g.filter_nodes(
              lambda v: (v.data["pos"] < step) & ~v.data["mask"].bool(),
              nids["dec"],
          )
          for i in range(self.decoder.N):
              pre_func, post_func = self.decoder.pre_func(
                  i, "qkv"
              ), self.decoder.post_func(i)
              nodes, edges = nodes_d, edges_dd
              self.update_graph(
                  g, edges, [(pre_func, nodes)], [(post_func, nodes)]
              )
              pre_q, pre_kv = self.decoder.pre_func(
                  i, "q", 1
              ), self.decoder.pre_func(i, "kv", 1)
              post_func = self.decoder.post_func(i, 1)
              nodes_e, nodes_d, edges = nids["enc"], nodes_d, edges_ed
              self.update_graph(
                  g,
                  edges,
                  [(pre_q, nodes_d), (pre_kv, nodes_e)],
                  [(post_func, nodes_d)],
              )
  
          frontiers = g.filter_nodes(
              lambda v: v.data["pos"] == step - 1, nids["dec"]
          )
          out = self.generator(g.ndata["x"][frontiers])
          batch_size = frontiers.shape[0] // k
          vocab_size = out.shape[-1]
          # Mask output for complete sequence
          one_hot = th.zeros(vocab_size).fill_(-1e9).to(device)
          one_hot[eos_id] = 0
          mask = g.ndata["mask"][frontiers].unsqueeze(-1).float()
          out = out * (1 - mask) + one_hot.unsqueeze(0) * mask
  
          if log_prob is None:
              log_prob, pos = out.view(batch_size, k, -1)[:, 0, :].topk(
                  k, dim=-1
              )
              eos = th.zeros(batch_size, k).byte()
          else:
              norm_old = eos.float().to(device) + (
                  1 - eos.float().to(device)
              ) * np.power((4.0 + step) / 6, alpha)
              norm_new = eos.float().to(device) + (
                  1 - eos.float().to(device)
              ) * np.power((5.0 + step) / 6, alpha)
              log_prob, pos = (
                  (
                      (
                          out.view(batch_size, k, -1)
                          + (log_prob * norm_old).unsqueeze(-1)
                      )
                      / norm_new.unsqueeze(-1)
                  )
                  .view(batch_size, -1)
                  .topk(k, dim=-1)
              )
  
          _y = y.view(batch_size * k, -1)
          y = th.zeros_like(_y)
          _eos = eos.clone()
          for i in range(batch_size):
              for j in range(k):
                  _j = pos[i, j].item() // vocab_size
                  token = pos[i, j].item() % vocab_size
                  y[i * k + j, :] = _y[i * k + _j, :]
                  y[i * k + j, step] = token
                  eos[i, j] = _eos[i, _j] | (token == eos_id)
  
          if eos.all():
              break
          else:
              g.ndata["mask"][nids["dec"]] = (
                  eos.unsqueeze(-1)
                  .repeat(1, 1, max_len)
                  .view(-1)
                  .to(device)
              )
      return y.view(batch_size, k, -1)[:, 0, :].tolist()
  
imports: from .config import *
  from .act import *
  from .attention import *
  from .viz import *
  from .layers import *
  from .functions import *
  from .embedding import *
  import threading
  import dgl.function
  import torch
  import torch.nn.init
label_desc: Use image node_id 5 for calling the Transformer obj's underlying member method code with example usage: obj.infer(graph, max_len, eos_id, k, alpha) and returns: y