code: def plot_multi_head_attention(
      data,
      uttid_list,
      attn_dict,
      outdir,
      suffix="png",
      savefn=savefig,
      ikey="input",
      iaxis=0,
      okey="output",
      oaxis=0,
      subsampling_factor=4,
  ):
      """Plot multi head attentions.
  
      :param dict data: utts info from json file
      :param List uttid_list: utterance IDs
      :param dict[str, torch.Tensor] attn_dict: multi head attention dict.
          values should be torch.Tensor (head, input_length, output_length)
      :param str outdir: dir to save fig
      :param str suffix: filename suffix including image type (e.g., png)
      :param savefn: function to save
      :param str ikey: key to access input
      :param int iaxis: dimension to access input
      :param str okey: key to access output
      :param int oaxis: dimension to access output
      :param subsampling_factor: subsampling factor in encoder
  
      """
      for name, att_ws in attn_dict.items():
          for idx, att_w in enumerate(att_ws):
              data_i = data[uttid_list[idx]]
              filename = "%s/%s.%s.%s" % (
                  outdir,
                  uttid_list[idx],
                  name,
                  suffix,
              )
              dec_len = (
                  int(data_i[okey][oaxis]["shape"][0]) + 1
              )  # +1 for <eos>
              enc_len = int(data_i[ikey][iaxis]["shape"][0])
              is_mt = "token" in data_i[ikey][iaxis].keys()
              # for ASR/ST
              if not is_mt:
                  enc_len //= subsampling_factor
              xtokens, ytokens = None, None
              if "encoder" in name:
                  att_w = att_w[:, :enc_len, :enc_len]
                  # for MT
                  if is_mt:
                      xtokens = data_i[ikey][iaxis]["token"].split()
                      ytokens = xtokens[:]
              elif "decoder" in name:
                  if "self" in name:
                      # self-attention
                      att_w = att_w[:, :dec_len, :dec_len]
                      if "token" in data_i[okey][oaxis].keys():
                          ytokens = data_i[okey][oaxis][
                              "token"
                          ].split() + ["<eos>"]
                          xtokens = ["<sos>"] + data_i[okey][oaxis][
                              "token"
                          ].split()
                  else:
                      # cross-attention
                      att_w = att_w[:, :dec_len, :enc_len]
                      if "token" in data_i[okey][oaxis].keys():
                          ytokens = data_i[okey][oaxis][
                              "token"
                          ].split() + ["<eos>"]
                      # for MT
                      if is_mt:
                          xtokens = data_i[ikey][iaxis]["token"].split()
              else:
                  logging.warning("unknown name for shaping attention")
              fig = _plot_and_save_attention(
                  att_w, filename, xtokens, ytokens
              )
              savefn(fig, filename)
  
file_path: espnet/espnet/nets/pytorch_backend/transformer/plot.py
imports: import logging
  import os
  import numpy
  from espnet.asr import asr_utils
label_desc: Use image node_id 3 for calling a global function with example usage: plot_multi_head_attention(data, uttid_list, attn_dict, outdir, suffix, savefn, ikey, iaxis, okey, oaxis, subsampling_factor) without return types