jax.nn.log_softmax

jax.nn.log_softmax(x, axis=-1, where=None, initial=None)

Log-Softmax function.

Computes the logarithm of the softmax function, which rescales elements to the range \([-\infty, 0)\).

\[\mathrm{log\_softmax}(x)_i = \log \left( \frac{\exp(x_i)}{\sum_j \exp(x_j)} \right)\]

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 log_softmax should be computed. Either an integer or a tuple of integers.

• where (Union[Array, ndarray, bool_, number, bool, int, float, complex, None]) – Elements to include in the log_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.

See also

softmax()