Numerically Stable Softmax

Problem

Compute the softmax of a vector x: softmax(x)_i = exp(x_i) / sum(exp(x_j)).

Subtract the max of x before exponentiating for numerical stability.

Hint

Subtract the max value from each element before exp() to prevent overflow.

case 1

in: [[1,2,3]]

expected: [0.09003057317038046,0.24472847105479764,0.6652409557748219]

case 2

in: [[0,0]]

expected: [0.5,0.5]

case 3

in: [[1000,1000,1000]]

expected: [0.3333333333333333,0.3333333333333333,0.3333333333333333]

case 4

in: [[0]]

expected: [1]

Starter codePython

import math
def softmax(x):
    # Your code here
    pass

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