Softmax

题面

Write a program that computes the softmax function for an array of 32-bit floating-point numbers on a GPU. For an input array x of length n, the softmax σ(x) is an array of length n where the i-th element is: σ(x)_i = e^{x_i} / Σ_j e^{x_j}. 为避免溢出，应使用 “max trick”（先减去最大值再做指数）。

Implementation Requirements

• Use only native features (external libraries are not permitted)
• The solve function signature must remain unchanged
• The final result must be stored in the array output

Examples

Example 1:

Input: [1.0, 2.0, 3.0], N = 3
Output: [0.090, 0.244, 0.665] (approximately)

Example 2:

Input: [-10.0, -5.0, 0.0, 5.0, 10.0], N = 5
Output: [2.04e-09, 4.52e-07, 9.99e-01, 2.26e-02, 9.77e-01] (approximately)

Constraints

• 1 ≤ N ≤ 500,000
• Performance is measured with N = 500,000