softmax(data, length=None, axis=-1, temperature=None, use_length=False, dtype=None)

Applies the softmax function.

The resulting array contains elements in the range (0,1) and the elements along the given axis sum up to 1.

\[softmax(\mathbf{z/t})_j = \frac{e^{z_j/t}}{\sum_{k=1}^K e^{z_k/t}}\]

for \(j = 1, ..., K\)

t is the temperature parameter in softmax function. By default, t equals 1.0

  • data (NDArray) – The input array.

  • axis (int, optional, default='-1') – The axis along which to compute softmax.

  • length (NDArray) – The length array.

  • temperature (double or None, optional, default=None) – Temperature parameter in softmax

  • dtype ({None, 'float16', 'float32', 'float64'},optional, default='None') – DType of the output in case this can’t be inferred. Defaults to the same as input’s dtype if not defined (dtype=None).

  • use_length (boolean or None, optional, default=0) – Whether to use the length input as a mask over the data input.


out – The output of this function.

Return type

NDArray or list of NDArrays


>>> data = np.ones((2, 3))
>>> npx.softmax(data, axis=0)
array([[0.5, 0.5, 0.5],
    [0.5, 0.5, 0.5]])
>>> npx.softmax(data, axis=1)
array([[0.33333334, 0.33333334, 0.33333334],
    [0.33333334, 0.33333334, 0.33333334]])