Softmax

Implement the softmax function over the last axis of an array.

For a vector x, softmax is

softmax(x)_i = exp(x_i) / sum_j exp(x_j)

Written that way it overflows: exp(1000) is inf, and inf / inf is nan. Softmax is shift invariant, meaning softmax(x) == softmax(x - c) for any constant c, so subtract the maximum before exponentiating.

Your function has to work for a single vector of shape (n,) and for a batch of shape (batch, n), normalising each row independently.

Function signature

def softmax(x):

Example

Input

softmax(np.array([1.0, 2.0, 3.0]))

Output

array([0.09003057, 0.24472847, 0.66524096])

Reasoning

Subtract the max, 3, to get [-2, -1, 0]. The exponentials are [0.135, 0.368, 1], with sum 1.503. Dividing gives the output.

Constraints and notes

• Subtract the maximum before exp.
• Work over the last axis, for shape (n,) or (batch, n).