Causal Attention With Temperature

Problem

Apply a causal mask and a temperature-scaled softmax to attention scores. Position i attends only to j <= i, and each kept score is divided by temperature before the stable softmax.

The signature is causal_attention_temperature(scores, temperature). Masked future entries become weight 0.0.

Hint

Each row has its own maximum over the unmasked prefix, and future columns stay exactly zero.

Test cases

case 1

in: [[[1,2],[3,4]],1]

expected: [[1,0],[0.2689414214,0.7310585786]]

case 2

in: [[[0,0],[0,0]],2]

expected: [[1,0],[0.5,0.5]]

case 3

in: [[[1]],0.5]

expected: [[1]]

case 4

in: [[[2,2,100],[0,1,0]],1]

expected: [[1,0,0],[0.2689414214,0.7310585786,0]]

Starter codePython

import math
def causal_attention_temperature(scores, temperature):
    # Your code here
    pass

4 test cases · run your code against them in the editor