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?
Thank you! We'll create more content like this.
Thank you for helping us improve!