code: def _argminmax(
      is_min: bool,
      operand: TfVal,
      axes: Sequence[int],
      index_dtype: DType,
      _in_avals: Sequence[core.ShapedArray],
      _out_aval: core.ShapedArray,
  ):
      # The following is known to diverge from JAX behavior for NaN.
      (axis,) = axes
      output_type = tf.int32
      if dtypes.iinfo(index_dtype).bits > 32:
          output_type = tf.int64
      # TODO(phawkins): handle axes larger than 2^31.
      fn = tf.math.argmin if is_min else tf.math.argmax
      result = fn(operand, axis=axis, output_type=output_type)
      return tf.cast(result, jax2tf._to_tf_dtype(index_dtype))
  
imports: from __future__ import annotations
  import builtins
  from collections.abc import Sequence
  import dataclasses
  from functools import partial, wraps
  import math
  import string
  from typing import Any, Callable, Optional
  from jax._src import core
  from jax import lax
  from jax._src.lax import slicing
  from jax._src import dtypes
  from jax._src import util
  from jax.experimental.jax2tf import jax2tf
  import numpy
  import tensorflow
label_desc: Use image node_id 18 for calling a global function with example usage: _argminmax(is_min, operand, axes, index_dtype, _in_avals, _out_aval) and returns: tf