Skip to main content

Mojo function

argmaxmin_gpu

argmaxmin_gpu[dtype: DType, output_type: DType, rank: Int, largest: Bool](ctx: DeviceContext, input: NDBuffer[dtype, rank, origin], output: NDBuffer[output_type, rank, origin])

Wraps the Top-K GPU kernel with K=1 to perform argmax on the inner-most dimension.

Parameters:

  • dtype (DType): DType - The data dtype of the input tensor.
  • output_type (DType): DType - The data dtype of the output tensor.
  • rank (Int): Int - The rank of the input tensor.
  • largest (Bool): Bool - Whether to perform argmax or argmin.

Was this page helpful?