jax.nn.softmax#

jax.nn.softmax(x, axis=-1, where=None, initial=None)[source]#

Softmax function.

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 (Union[Array, ndarray, bool_, number, bool, int, float, complex]) – input array
axis (Union[int, tuple[int, ...], None]) – the axis or axes along which the softmax should be computed. The softmax output summed across these dimensions should sum to \(1\). Either an integer or a tuple of integers.
where (Union[Array, ndarray, bool_, number, bool, int, float, complex, None]) – Elements to include in the softmax.
initial (Union[Array, ndarray, bool_, number, bool, int, float, complex, None]) – The minimum value used to shift the input array. Must be present when where is not None.

Return type:

Array

Returns:

An array.

jax.nn.softmax