triton.language.argmin

triton.language.argmin(input, axis)

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

Parameters:
  • input – the input values

  • axis – the dimension along which the reduction should be done