code: def _softmax_jvp(axis, primals, tangents):
      (x, where, initial), (x_dot, _, _) = primals, tangents
      y = _softmax(x, axis, where, initial)
      return y, y * (
          x_dot - (y * x_dot).sum(axis, where=where, keepdims=True)
      )
  
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 18 for calling a global function with example usage: _softmax_jvp(axis, primals, tangents) and returns: y,  unknown