triton.language.argmin

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

返回输入 input tensor 中沿指定 axis 的最小元素的索引

参数:

• input (Tensor) – 输入值

• axis (int) – 执行归约操作的维度。如果为 None，则归约所有维度

• keep_dims (bool) – 如果为 true，则保留长度为 1 的归约维度

• tie_break_left (bool) – 如果为 true，在出现平局时（即多个元素具有相同的最小索引值），返回非 NaN 值的最左侧索引

此函数也可以作为成员函数在 tensor 上调用，例如 x.argmin(...) 而不是 argmin(x, ...)。