Implement a causal attention head

Intermediate65 min

Implement causal_attention(queries, keys, values) for T positions. Queries and keys are T-by-D finite numeric matrices, values are T-by-V, and D and V are positive. Output a T-by-V matrix. Row i attends only to positions 0 through i using scores divided by sqrt(D). Empty matrices together return []. Reject incompatible or ragged shapes and nonfinite entries. Use a stable softmax and never normalize over masked positions.

Your task

Complete the starter function using the contract above.
Use the examples and visible tests to check normal inputs, boundaries, and rejected inputs.
Run tests to record your result, then compare with the explained reference solution.

Examples

EXAMPLE 1

Inputcausal_attention([[0],[0]], [[1],[2]], [[2],[6]])

Output[[2.0], [4.0]]

Equal visible scores give a prefix average.

EXAMPLE 2

Inputcausal_attention([[1]], [[2]], [[3,4]])

Output[[3.0,4.0]]

One visible value is returned exactly regardless of its score.

def causal_attention(queries, keys, values):
    raise NotImplementedError