code: def log_softmax(
      x: ArrayLike,
      axis: int | tuple[int, ...] | None = -1,
      where: ArrayLike | None = None,
      initial: ArrayLike | None = None,
  ) -> Array:
      r"""Log-Softmax function.
  
      Computes the logarithm of the :code:`softmax` function, which rescales
      elements to the range :math:`[-\infty, 0)`.
  
      .. math ::
        \mathrm{log\_softmax}(x)_i = \log \left( \frac{\exp(x_i)}{\sum_j \exp(x_j)}
        \right)
  
      Args:
        x : input array
        axis: the axis or axes along which the :code:`log_softmax` should be
          computed. Either an integer or a tuple of integers.
        where: Elements to include in the :code:`log_softmax`.
        initial: The minimum value used to shift the input array. Must be present
          when :code:`where` is not None.
  
      Returns:
        An array.
  
      See also:
        :func:`softmax`
      """
      numpy_util.check_arraylike("log_softmax", x)
      x_arr = jnp.asarray(x)
      x_max = jnp.max(
          x_arr, axis, where=where, initial=initial, keepdims=True
      )
      shifted = x_arr - lax.stop_gradient(x_max)
      shifted_logsumexp = jnp.log(
          jnp.sum(jnp.exp(shifted), axis, where=where, keepdims=True)
      )
      result = shifted - shifted_logsumexp
      if where is not None:
          return jnp.where(where, result, -jnp.inf)
      return result
  
imports: from __future__ import annotations
  from functools import partial
  import operator
  import numpy
  from typing import Any
  import jax
  import jax.numpy
  from jax import custom_jvp
  from jax import lax
  from jax._src import config
  from jax._src import core
  from jax._src import dtypes
  from jax._src import util
  from jax._src.core import AxisName
  from jax._src.numpy import util
  from jax._src.typing import Array, ArrayLike
  from jax._src.ops.special import logsumexp
label_desc: Use image node_id 15 for calling a global function with example usage: log_softmax(x, axis, where, initial) and returns: result, jnp