code: def argmax(x, /, *, axis=None, keepdims=False):
      """Returns the indices of the maximum values along a specified axis."""
      return jax.numpy.argmax(x, axis=axis, keepdims=keepdims)
  
imports: import jax
  from jax.experimental.array_api._data_type_functions import result_type
label_desc: Use image node_id 1 for calling a global function with example usage: argmax() and returns: jax