Skip to main content
Log in

Mojo function

arg_max

arg_max(v: Symbol, axis: Int = -1) -> Symbol

Finds the index of the maximum value along a dimension.

Args:

  • v (Symbol): The input tensor for the operation.
  • axis (Int): The axis along which to compute the reduction. If negative, indexes from the last dimension, ie. a value of -1 will compute the reduction along the last dimension.

Returns:

A symbolic tensor representing the result of the arg_max operation. The tensor will have the same rank as the input tensor, and the same shape except along the axis dimension which will have size 1.