keras/src/ops/nn.py/log_softmax
def log_softmax(x, axis=-1):
    """Log-softmax activation function.

    It is defined as:
    `f(x) = x - max(x) - log(sum(exp(x - max(x))))`

    Args:
        x: Input tensor.
        axis: Integer, axis along which the log-softmax is applied.
            Defaults to `-1`.

    Returns:
        A tensor with the same shape as `x`.

    Example:

    >>> x = np.array([-1., 0., 1.])
    >>> x_log_softmax = keras.ops.log_softmax(x)
    >>> print(x_log_softmax)
    array([-2.40760596, -1.40760596, -0.40760596], shape=(3,), dtype=float64)

    """
    if any_symbolic_tensors((x,)):
        return LogSoftmax(axis).symbolic_call(x)
    if isinstance(axis, tuple):
        original_shape = x.shape
        new_shape = []
        skip_dims = set(axis)
        i = 0
        while i < len(original_shape):
            if i in skip_dims:
                size = 1
                while i in skip_dims:
                    size *= original_shape[i]
                    i += 1
                new_shape.append(size)
            else:
                new_shape.append(original_shape[i])
                i += 1
        x = backend.numpy.reshape(x, new_shape)
        x = backend.nn.log_softmax(x, axis=-1)
        x = backend.numpy.reshape(x, original_shape)
        return x
    else:
        return backend.nn.log_softmax(x, axis=axis)