def _get_attention(tokens, attn):
    """Compute representation of the attention to pass to the d3 visualization

    Args:
      tokens: tokens in input text
      attn: numpy array, attention
          [num_layers, batch_size, num_heads, seq_len, seq_len]

    Returns:
      Dictionary:
      {
        'att': list of inter attentions matrices, one for each layer. Each is of shape [num_heads, source_seq_len, target_seq_len]
        'left_text': list of source tokens, to be displayed on the left of the vis
        'right_text': list of target tokens, to be displayed on the right of the vis
      }
    """

    all_attns = []
    num_layers = len(attn)
    for layer in range(num_layers):
        layer_attn = attn[layer][0] # Get layer attention (assume batch size = 1), shape = [num_heads, seq_len, seq_len]
        all_attns.append(layer_attn.tolist()) # Append attention for layer, across all heads

    return {
        'att': all_attns,
        'left_text': tokens,
        'right_text': tokens
    }