jax.nn.softmax#
- jax.nn.softmax(x, axis=-1, where=None, initial=_UNSPECIFIED)[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 (ArrayLike) – input array
axis (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 (ArrayLike | None) – Elements to include in the
softmax
.
- Return type:
Array
- Returns:
An array.
See also
- Parameters:
initial (ArrayLike | None | Unspecified)