code: def softmax(
      x: ArrayLike,
      axis: int | tuple[int, ...] | None = -1,
      where: ArrayLike | None = None,
      initial: ArrayLike | None = None,
  ) -> Array:
      r"""Softmax function.
  
      Computes the function which rescales elements to the range :math:`[0, 1]`
      such that the elements along :code:`axis` sum to :math:`1`.
  
      .. math ::
        \mathrm{softmax}(x) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
  
      Args:
        x : input array
        axis: the axis or axes along which the softmax should be computed. The
          softmax output summed across these dimensions should sum to :math:`1`.
          Either an integer or a tuple of integers.
        where: Elements to include in the :code:`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:`log_softmax`
      """
      if config.softmax_custom_jvp.value:
          # mypy is confused by the `functools.partial` application in the definition
          # of `_softmax` and incorrectly concludes that `_softmax` returns
          # `ReturnValue` -- the unsubstituted type parameter of `custom_jvp`.
          return _softmax(x, axis, where, initial)  # type: ignore[return-value]
      else:
          return _softmax_deprecated(x, axis, where, initial)
  
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 16 for calling a global function with example usage: softmax(x, axis, where, initial) and returns: _softmax, _softmax_deprecated