code: def categorical(key, logits, axis, shape=None):
      logits_arr = ivy.asarray(logits)
  
      if axis >= 0:
          axis -= len(logits_arr.shape)
      batch_shape = tuple(_remove_axis(logits_arr.shape, axis))
  
      if shape is None:
          shape = batch_shape
      else:
          shape = tuple(shape)
          if shape != batch_shape:
              raise ValueError(
                  +f"Shape {shape} is not compatible with reference shape {batch_shape}"
              )
  
      logits_shape = list(shape[len(shape) - len(batch_shape) :])
      logits_shape.insert(
          axis % len(logits_arr.shape), logits_arr.shape[axis]
      )
  
      gumbel_noise = gumbel(
          key, ivy.array(logits_shape), logits_arr.dtype
      )
      expanded_logits = ivy.expand_dims(logits_arr, axis=axis)
      noisy_logits = gumbel_noise + expanded_logits
  
      # Use Ivy's argmax to get indices
      indices = ivy.argmax(noisy_logits, axis=axis)
  
      return indices
  
imports: import operator
  import ivy
  from ivy.func_wrapper import with_unsupported_dtypes, with_supported_dtypes
  from ivy.functional.frontends.jax.func_wrapper import to_ivy_arrays_and_back, handle_jax_dtype
label_desc: Use image node_id 7 for calling a global function with example usage: categorical(key, logits, axis, shape) and returns: indices