jax.scipy.special.softmax#

jax.scipy.special.softmax(x, /, *, axis=None)[source]#

Softmax function.

JAX implementation of scipy.special.softmax().

Computes the function which rescales elements to the range \([0, 1]\) such that the elements along axis sum to \(1\).

\[\mathrm{softmax}(x) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}\]
Parameters:

  • x (ArrayLike) – input array

  • axis (int | tuple[int, ...] | None | None) – the axis or axes along which the softmax should be computed. The softmax output summed across these dimensions should sum to \(1\).

Returns:

An array of the same shape as x.

Return type:

Array

Note

If any input values are +inf, the result will be all NaN: this reflects the fact that inf / inf is not well-defined in the context of floating-point math.

See also

log_softmax()