Easy Primitives

Implement Softmax

Softmax turns a vector of scores, often called logits, into probabilities. Each score is exponentiated and divided by the sum of all those exponentials:

softmax(x)[i] = exp(x[i]) / sum(exp(x))

For [1.0, 2.0, 3.0], the result is approximately [0.090031, 0.244728, 0.665241]. The largest input gets the largest probability. Equal inputs get equal probabilities; four equal scores each get 0.25.

Subtract the maximum first

Direct exponentiation can overflow for large positive scores. Subtracting the same constant from every score leaves the mathematical probabilities unchanged. Choose the maximum so the largest shifted score is zero:

shifted = x - max(x)
weights = exp(shifted)
probabilities = weights / sum(weights)

For [1000.0, 1001.0, 1002.0], the shifted scores are [-2.0, -1.0, 0.0]. Their largest exponential is one, so this avoids the overflow in exp(1002). The probabilities match those for [1.0, 2.0, 3.0], up to floating-point error.

Your task

Implement softmax(x) for a nonempty one-dimensional tensor of finite floating-point scores. Return a vector of the same shape, normalized across the whole vector. Use basic reductions and arithmetic; do not call torch.softmax, torch.nn.functional.softmax or jax.nn.softmax.

The result should sum to approximately one. Extremely small probabilities can round to zero, so do not assume every floating-point result is strictly positive. NaN and infinite inputs are outside this exercise's contract.

The explorer below lets you change temperature. This exercise uses the ordinary formula at temperature one, with no additional temperature argument.

Loading visualization…

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Turn scores into probabilities

Softmax converts a list of scores, also called logits, into probabilities that sum to one. It exponentiates each score and divides by the sum:

$$p_i=\frac{\exp(z_i/T)}{\sum_j\exp(z_j/T)},\qquad T>0.$$

T is the temperature. With the usual choice T = 1, scores [0, 0] give probabilities [0.5, 0.5]. Scores [0, log(3)] give [0.25, 0.75], because their exponentials are [1, 3].

Lower temperature puts more probability on the largest scores. Higher temperature makes the probabilities more even. A positive temperature preserves the ordering of finite scores.

Subtract the maximum before exponentiating

Large positive scores can make exponentials overflow. Subtracting the largest scaled score avoids large positive exponents without changing the probabilities: the common exponential factor cancels between numerator and denominator.

For a finite, nonempty list of scaled scores, every shifted score is at most zero and at least one is exactly zero. Their exponentials are therefore at most one, with at least one equal to one. For example, [1000, 1000] becomes [0, 0] before exponentiation and still gives [0.5, 0.5].

Choose the axis that contains the alternatives

For logits shaped (batch, classes), apply softmax over classes. Each example then gets its own probability distribution. Summing over the entire tensor would incorrectly make examples compete with each other.

When implementing the division yourself, keep the summed axis as a size-one dimension so the denominator broadcasts back across the intended scores.

If you need log probabilities, use log-softmax or a log-sum-exp formulation. Taking the logarithm after softmax can produce negative infinity when a small probability has already rounded to zero.

For masked softmax, each row needs at least one allowed, finite score. A row containing only negative infinity has no valid choices to normalise.

Detecting runtime…

Implement Softmax

Subtract the maximum first

Your task

Turn scores into probabilities

Subtract the maximum before exponentiating

Choose the axis that contains the alternatives

Your code, saved here

Stuck?

JAX reference solution