Implement the causal information flow of one attention head as a prefix average: each position can read itself and the past, never the future.

def causal_attention(values):
    return [prefix_average, ...]

• include the current position in its receptive field
• exclude every future position
• normalize each prefix by its length