code: def resize(att_mat, max_length=None):
      """Normalize attention matrices and reshape as necessary."""
      for i, att in enumerate(att_mat):
          # Add extra batch dim for viz code to work.
          if att.ndim == 3:
              att = np.expand_dims(att, axis=0)
          if max_length is not None:
              # Sum across different attention values for each token.
              att = att[:, :, :max_length, :max_length]
              row_sums = np.sum(att, axis=2)
              # Normalize
              att /= row_sums[:, :, np.newaxis]
          att_mat[i] = att
      return att_mat
  
file_path: tensor2tensor/tensor2tensor/visualization/attention.py
imports: import json
  import os
  import IPython.display
  import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: resize(att_mat, max_length) and returns: att_mat