code: def draw_heatmap(array, input_seq, output_seq, dirname, name):
      dirname = os.path.join("log", dirname)
      if not os.path.exists(dirname):
          os.makedirs(dirname)
  
      fig, axes = plt.subplots(2, 4)
      cnt = 0
      for i in range(2):
          for j in range(4):
              axes[i, j].imshow(array[cnt].transpose(-1, -2))
              axes[i, j].set_yticks(np.arange(len(input_seq)))
              axes[i, j].set_xticks(np.arange(len(output_seq)))
              axes[i, j].set_yticklabels(input_seq, fontsize=4)
              axes[i, j].set_xticklabels(output_seq, fontsize=4)
              axes[i, j].set_title("head_{}".format(cnt), fontsize=10)
              plt.setp(
                  axes[i, j].get_xticklabels(),
                  rotation=45,
                  ha="right",
                  rotation_mode="anchor",
              )
              cnt += 1
  
      fig.suptitle(name, fontsize=12)
      plt.tight_layout()
      plt.savefig(os.path.join(dirname, "{}.pdf".format(name)))
      plt.close()
  
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 2 for calling a global function with example usage: draw_heatmap(array, input_seq, output_seq, dirname, name) without return types