def update_decoder_attention_history(cache):
      """Save attention weights in cache, e.g., for vizualization."""
      for k in [x for x in self.attention_weights
                if "decoder" in x and "self" not in x and "logits" not in x]:
        idx = k.find("layer_")
        if idx < 0:
          continue
        # Get layer number from the string name.
        layer_nbr = k[idx + 6:]
        idx = 0
        while idx + 1 < len(layer_nbr) and layer_nbr[:idx + 1].isdigit():
          idx += 1
        layer_nbr = "layer_%d" % int(layer_nbr[:idx])
        if layer_nbr in cache["attention_history"]:
          cache["attention_history"][layer_nbr] = tf.concat(
              [cache["attention_history"][layer_nbr],
               self.attention_weights[k]],
              axis=2)
Save attention weights in cache, e.g., for vizualization.
tensor2tensor/models/transformer.py