Row-Wise Softmax

Problem

Apply softmax to every row of a matrix independently.

For each row, subtract the row maximum before exponentiating for numerical stability, then divide by the sum of exponentials. Empty rows stay empty and a masked entry of -1000000000.0 becomes probability 0.0.

Hint

Converting scores to probabilities row by row is what makes attention a weighted average.

Test cases

case 1

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

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

case 2

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

expected: [[0.09003057317038046,0.24472847105479764,0.6652409557748218]]

case 3

in: [[[0.7071067811865475,-1000000000]]]

expected: [[1,0]]

case 4

in: [[[0]]]

expected: [[1]]

Starter codePython

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

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