code: def _softmax(
      x: ArrayLike,
      axis: int | tuple[int, ...] | None = -1,
      where: ArrayLike | None = None,
      initial: ArrayLike | None = None,
  ) -> Array:
      x_max = jnp.max(
          x, axis, where=where, initial=initial, keepdims=True
      )
      unnormalized = jnp.exp(x - x_max)
      result = unnormalized / jnp.sum(
          unnormalized, axis, where=where, keepdims=True
      )
      if where is not None:
          result = jnp.where(where, result, 0)
      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 17 for calling a global function with example usage: _softmax(x, axis, where, initial) and returns: result