code: def draw_atts(maps, src, tgt, dirname, prefix):
      """
      maps[0]: encoder self-attention
      maps[1]: encoder-decoder attention
      maps[2]: decoder self-attention
      """
      draw_heatmap(
          maps[0], src, src, dirname, "{}_enc_self_attn".format(prefix)
      )
      draw_heatmap(
          maps[1], src, tgt, dirname, "{}_enc_dec_attn".format(prefix)
      )
      draw_heatmap(
          maps[2], tgt, tgt, dirname, "{}_dec_self_attn".format(prefix)
      )
  
file_path: dgl/examples/pytorch/transformer/modules/viz.py
imports: import os
  import matplotlib
  import matplotlib.animation
  import matplotlib.pyplot
  import networkx
  import numpy
  import torch
  from networkx.algorithms import bipartite
  import networkx
  from matplotlib.patches import ConnectionStyle, FancyArrowPatch
  from networkx.utils import is_string_like
label_desc: Use image node_id 3 for calling a global function with example usage: draw_atts(maps, src, tgt, dirname, prefix) without return types