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.