code: def argmax(
      x: JaxArray,
      /,
      *,
      axis: Optional[int] = None,
      keepdims: bool = False,
      dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None,
      select_last_index: bool = False,
      out: Optional[JaxArray] = None,
  ) -> JaxArray:
      if select_last_index:
          x = jnp.flip(x, axis=axis)
          ret = jnp.array(jnp.argmax(x, axis=axis, keepdims=keepdims))
          if axis is not None:
              ret = x.shape[axis] - ret - 1
          else:
              ret = x.size - ret - 1
      else:
          ret = jnp.argmax(x, axis=axis, keepdims=keepdims)
      if dtype:
          dtype = ivy.as_native_dtype(dtype)
          return ret.astype(dtype)
      return ret
  
file_path: ivy/ivy/functional/backends/jax/searching.py
imports: from numbers import Number
  from typing import Optional, Tuple, Union
  import jax.numpy
  import ivy
  from ivy.functional.backends.jax import JaxArray
  from .None import backend_version
  from ivy.func_wrapper import with_unsupported_dtypes
label_desc: Use image node_id 1 for calling a global function with example usage: argmax() and returns: ret, ret