jax.nn.log_softmax#

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

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 (ArrayLike) – input array

  • axis (Axis) – the axis or axes along which the log_softmax should be computed. Either an integer or a tuple of integers.

  • where (ArrayLike | None) – Elements to include in the log_softmax. The output for any masked-out element is minus infinity.

Returns:

An array.

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

softmax()