roofline.devPublic Beta

Argmax Over Dimension

0EASYRankedReport issue

Given a tensor $A$ with rank $n$ and shape $\text{shape}$ stored in row-major order, compute the maximum value over a specified axis $\text{dim}$ .

B_{i_0 \, \dots \, i_{\text{dim}-1} \, i_{\text{dim}+1} \, \dots \, i_{n-1}} = \max_{i_\text{dim}} A_{i_0 \, \dots \, i_{n-1}}

Input

A - input tensor stored in row-major order containing single-precision floating-point values.
shape - integer array describing the dimensions of the input tensor; its length is n.
n - the number of dimensions (rank) of the tensor.
dim - the axis along which to reduce (0-indexed, valid in [0, n)).

Output

B - tensor containing the maximum values from reducing A along dimension dim.

Open on a desktop browser to write and submit code.