warp.tile_argmax#

warp.tile_argmax(
a: Tile[Scalar, tuple[int, ...]],
) → Tile[Int, tuple[Literal[1]]]#
  • Kernel: true
  • Python: false
  • Differentiable: false

Compute the index of the maximum element of a tile.

This is a cooperative operation: the threads of the block divide the work between them, so every thread must reach the call.

The index is the row-major linear index into the flattened tile, so in a shape-(2, 4) tile the element at row 1, column 1 has index 5. If several elements are equal maxima, which of their indices is returned is unspecified: it depends on block_dim and on the device.

Currently, if a contains NaN, the returned index may point to a NaN element and may not match tile_max().

Parameters:

a – The tile to reduce. Must have a scalar element type.

Returns:

A shape-(1,) register tile of warp.int32 holding the index. Read the value with [0].

See tile_argmin() for an analogous usage example.