argmax¶

nvtripy.argmax(input: Tensor, dim: int | None = None, keepdim: bool = False) → Tensor[source]¶

Returns a new tensor containing the indices of maximum values of the input tensor along the specified dimension.

If there are multiple maximum values, then the indices of the first maximum value are returned.

Parameters:

input (Tensor) – The input tensor.
dim (int | None) – The dimension along which to reduce. If this is not provided, the index of the flattened input is returned.
keepdim (bool) – Whether to retain reduced dimensions in the output. If this is False, reduced dimensions will be squeezed.

Returns:

A new tensor.

Return type:

Tensor

INPUT REQUIREMENTS:: input.dtype is one of [float32, float16, bfloat16, int32, int64]
OUTPUT GUARANTEES:: return[0].dtype == int32

Example

input = tp.Tensor([[1.0, 0.0, 3.0], [0.5, 2.0, 1.5]])
output = tp.argmax(input, 0)

Local Variables¶

>>> input
tensor(
    [[1, 0, 3],
     [0.5, 2, 1.5]], 
    dtype=float32, loc=cpu:0, shape=(2, 3))

>>> output
tensor([0, 1, 0], dtype=int32, loc=gpu:0, shape=(3,))