Compute the indices of maximum elements along the given axis.

The argument axis could be negative or None:

x = dragon.constant([[1, 2, 3], [4, 5, 6]])

# A negative ``axis`` is the last-k axis
print(dragon.math.argmax(x, 1))
print(dragon.math.argmax(x, -1))  # Equivalent

# If ``axis`` is None, the vector-style reduction
# will be applied to return a scalar index
print(dragon.math.argmax(x))  # 5
  • inputs (dragon.Tensor) – The input tensor.
  • axis (int, optional) – The axis to reduce.
  • top_k (int, optional, default=1) – The top k results to keep.
  • keep_dims (bool, optional, default=False) – Keep the reduced dimension or not.

dragon.Tensor – The indices of elements.