Stable Softmax Rows

Problem

Apply the numerically stable softmax to each row of matrix A independently.

For every row, subtract the row maximum before exponentiating, then normalize so each output row sums to 1.

Hint

Max subtraction per row prevents overflow when one entry is huge.

case 1

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

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

case 2

in: [[[0,1]]]

expected: [[0.2689414213699951,0.7310585786300049]]

case 3

in: [[[0,1000]]]

expected: [[0,1]]

case 4

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

expected: [[0.09003057317038046,0.24472847105479764,0.6652409557748218]]

Starter codePython

def softmax_rows(A):
    # Your code here
    pass

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