code: def _approx_top_k_jvp(
      primals,
      tangents,
      *,
      k,
      reduction_dimension,
      recall_target,
      is_max_k,
      reduction_input_size_override,
      aggregate_to_topk,
  ):
      (operand,) = primals
      (tangent,) = tangents
      if is_max_k:
          val_out, arg_out = approx_max_k(
              operand,
              k,
              reduction_dimension,
              recall_target,
              reduction_input_size_override,
              aggregate_to_topk,
          )
      else:
          val_out, arg_out = approx_min_k(
              operand,
              k,
              reduction_dimension,
              recall_target,
              reduction_input_size_override,
              aggregate_to_topk,
          )
      if type(tangent) is ad_util.Zero:
          tangent_out = ad_util.Zero.from_value(val_out)
      else:
          arg_shape = arg_out.shape
          rank = len(arg_shape)
          if reduction_dimension < 0:
              reduction_dimension += rank
          iotas = [
              lax.broadcasted_iota(arg_out.dtype, arg_shape, i)
              for i in range(rank)
          ]
          idx = tuple(
              arg_out if i == reduction_dimension else iotas[i]
              for i in range(rank)
          )
          tangent_out = tangent[idx]
      return (val_out, arg_out), (
          tangent_out,
          ad_util.Zero.from_value(arg_out),
      )
  
imports: from functools import partial
  import numpy
  from jax._src import ad_util
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.interpreters import xla
  from jax._src.lax import lax
  from jax._src.lib import xla_client
  from jax._src.lib.mlir import ir
  from jax._src.lib.mlir.dialects import func
  from jax._src.lib.mlir.dialects import hlo
  from jax._src.typing import Array
label_desc: Use image node_id 10 for calling a global function with example usage: _approx_top_k_jvp(primals, tangents) and returns: val_out,  arg_out,  tangent_out,  ad_util