triton.language.argmin#

triton.language.argmin(input, axis, tie_break_left=True, keep_dims=False)#

Returns the minimum index of all elements in the input tensor along the provided axis

Parameters:

input (Tensor) – the input values
axis (int) – the dimension along which the reduction should be done. If None, reduce all dimensions
keep_dims (bool) – if true, keep the reduced dimensions with length 1
tie_break_left (bool) – if true, in case of a tie (i.e., multiple elements have the same minimum index value), return the left-most index for values that aren’t NaN

This function can also be called as a member function on tensor, as x.argmin(...) instead of argmin(x, ...).

Example

import torch
import torch_npu
import triton
import triton.language as tl


@triton.jit
def triton_argmin_2d(in_ptr0, out_ptr0, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr,
                     NNUMEL: tl.constexpr):
    mblk_idx = tl.arange(0, MNUMEL)
    nblk_idx = tl.arange(0, NNUMEL)
    mmask = mblk_idx < M
    nmask = nblk_idx < N
    mask = (mmask[:, None]) & (nmask[None, :])
    idx = mblk_idx[:, None] * N + nblk_idx[None, :]
    x = tl.load(in_ptr0 + idx, mask=mask, other=float('inf'))
    tmp4 = tl.argmin(x, dim)
    if dim == 0:
        tl.store(out_ptr0 + tl.arange(0, N), tmp4, None)
    else:
        tl.store(out_ptr0 + tl.arange(0, M), tmp4, None)


def test_argmin():
    M, N, dim = 4, 8, 1
    x = torch.randn(M, N, dtype=torch.float32).npu()
    out = torch.empty(M, dtype=torch.int32).npu()
    triton_argmin_2d[1, 1, 1](x, out, dim, M, N, M, N)
    ref = torch.argmin(x, dim=dim).to(torch.int32)
    assert torch.equal(out.cpu(), ref.cpu()), "argmin result mismatch"
    print("argmin result:", out)


if __name__ == "__main__":
    test_argmin()

DataType Support

平台	uint8	int8	uint16	int16	uint32	int32	uint64	int64	fp16	fp32	fp64	bf16	fp8e(e4m3)	fp8e5(e5m2)	bool
Ascend A2/A3	√	√	×	√	×	√	×	√	√	√	×	√	×	×	√
Ascend 950	√	√	√	√	√	√	√	√	√	√	×	√	×	×	√

Special Restrictions

keep_dims=True requires more test coverage; currently verified for 3D tensor with dim=2.

triton.language.argmin

Contents

triton.language.argmin#