argmin¶

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

Returns a new tensor containing the indices of minimum values of the input tensor along the specified dimension. If there are multiple minimum values, then the indices of the first minimum 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.argmin(input, 0)

Local Variables¶

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

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