Softmax

Given an array x of N floats, write its softmax into y:

y[i] = exp(x[i]) / (exp(x[0]) + exp(x[1]) + ... + exp(x[N-1]))

Every y[i] lies in [0, 1] and they add up to 1. In float32, exp overflows to infinity once its input exceeds about 88.7, and infinity divided by infinity is NaN. Since y[i] does not change when the same constant is subtracted from every x[i], subtract the maximum m of x first:

y[i] = exp(x[i] - m) / (exp(x[0] - m) + ... + exp(x[N-1] - m))

Now every exponent is at most 0, so nothing overflows.

Requirements

• Implement solve in the starter of the framework you choose; keep its signature.
  • CUDA: extern "C" void solve(const float* x, float* y, int N)
  • Triton, PyTorch: def solve(x: torch.Tensor, y: torch.Tensor, N: int)
• x and y are already on the GPU. Do not modify x.
• Write the result into y in place; its contents on entry are undefined.
• You may allocate temporary GPU memory inside solve.

Examples

Input:  x = [0.0, 1.0, 2.0, 3.0], N = 4
Output: y = [0.032059, 0.087144, 0.236883, 0.643914]

Input:  x = [100.0, 100.0, 99.0], N = 3
Output: y = [0.422319, 0.422319, 0.155362]

In the second example exp(100.0) is already infinity in float32; after subtracting m = 100.0 the exponents are 0, 0 and -1.

Constraints

• 1 ≤ N ≤ 500,000
• -100 ≤ x[i] ≤ 100
• An output is accepted if it is within 1e-9 + 2e-3 × |expected| of the expected value
• Performance is measured with N = 500,000