Attention Softmax Weights

Problem

Turn attention score rows into probability distributions with a numerically stable row-wise softmax: subtract each row's maximum before exponentiating.

The signature is attention_weights(scores). Each output row sums to 1.0.

Hint

Normalize independently per row; the max subtraction keeps large scores from overflowing.

Test cases

case 1

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

expected: [[0.0900305732,0.2447284711,0.6652409558]]

case 2

in: [[[0,0],[1,1]]]

expected: [[0.5,0.5],[0.5,0.5]]

case 3

in: [[[-1,0,1]]]

expected: [[0.0900305732,0.2447284711,0.6652409558]]

case 4

in: [[[0]]]

expected: [[1]]

case 5

in: [[[1000,1000]]]

expected: [[0.5,0.5]]

Starter codePython

import math
def attention_weights(scores):
    # Your code here
    pass

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