code: def get_att_mats(translate_model):
    
    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
docs: 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)
    )