argmax

→Tensor

argmax(x: Tensor, dim: _int | None = None, keepdim: _bool = False)

source

Return the indices of the maximum values along dim.

For each slice along dim the index of the largest element is returned. When dim is None the tensor is treated as flattened and a single scalar index into the flattened layout is produced. Ties are broken by returning the first occurrence in row-major order.

Parameters

xTensor

Input tensor.

dimint= None

Dimension along which to reduce. None means reduce over the whole tensor.

keepdimbool= False

Retain the reduced dimension with size 1 when True.

Returns

Tensor

int64 tensor of indices. Shape matches x.shape with dim removed (or set to 1 when keepdim=True).

Notes

Mathematically, for a 1-D input $x$ ,

\mathrm{argmax}(x) = \underset{i}{\arg\max}\; x_i.

Examples

>>> import lucid
>>> x = lucid.tensor([[1.0, 5.0, 2.0], [4.0, 0.0, 3.0]])
>>> lucid.argmax(x, dim=1)
Tensor([1, 0])