code: def get_att_mats(translate_model):
      """Get's the tensors representing the attentions from a build model.
  
      The attentions are stored in a dict on the Transformer object while building
      the graph.
  
      Args:
        translate_model: Transformer object to fetch the attention weights from.
  
      Returns:
      Tuple of attention matrices; (
          enc_atts: Encoder self attention weights.
            A list of `num_layers` numpy arrays of size
            (batch_size, num_heads, inp_len, inp_len)
          dec_atts: Decoder self attetnion weights.
            A list of `num_layers` numpy arrays of size
            (batch_size, num_heads, out_len, out_len)
          encdec_atts: Encoder-Decoder attention weights.
            A list of `num_layers` numpy arrays of size
            (batch_size, num_heads, out_len, inp_len)
      )
      """
      enc_atts = []
      dec_atts = []
      encdec_atts = []
  
      prefix = "transformer/body/"
      postfix_self_attention = (
          "/multihead_attention/dot_product_attention"
      )
      if (
          translate_model.hparams.self_attention_type
          == "dot_product_relative"
      ):
          postfix_self_attention = (
              "/multihead_attention/" "dot_product_attention_relative"
          )
      postfix_encdec = "/multihead_attention/dot_product_attention"
  
      for i in range(translate_model.hparams.num_hidden_layers):
          enc_att = translate_model.attention_weights[
              "%sencoder/layer_%i/self_attention%s"
              % (prefix, i, postfix_self_attention)
          ]
          dec_att = translate_model.attention_weights[
              "%sdecoder/layer_%i/self_attention%s"
              % (prefix, i, postfix_self_attention)
          ]
          encdec_att = translate_model.attention_weights[
              "%sdecoder/layer_%i/encdec_attention%s"
              % (prefix, i, postfix_encdec)
          ]
          enc_atts.append(enc_att)
          dec_atts.append(dec_att)
          encdec_atts.append(encdec_att)
  
      return enc_atts, dec_atts, encdec_atts
  
file_path: tensor2tensor/tensor2tensor/visualization/visualization.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from tensor2tensor import models
  from tensor2tensor import problems
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import trainer_lib
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 2 for calling a global function with example usage: get_att_mats(translate_model) and returns: enc_atts,  dec_atts,  encdec_atts