import torch


def format_attn(attention_tuples: tuple):
    """
    Input: N tuples (N = layer num)

    Each tuple item is Tensor of shape
    Batch x num heads x from x to

    Output: Tensor of shape layer x from x to
    (averaged over heads)
    """

    # Combine tuples into large Tensor, then avg
    return torch.cat([l for l in attention_tuples], dim=0).mean(dim=1)