Skip to main content

Mojo function

argsort

argsort[*, ascending: Bool = True, target: StringSlice[StaticConstantOrigin] = StringSlice("cpu")](output: TileTensor[output.dtype, output.LayoutType, output.origin, linear_idx_type=output.linear_idx_type, element_size=output.element_size], input: TileTensor[input.dtype, input.LayoutType, input.origin, address_space=input.address_space, linear_idx_type=input.linear_idx_type, element_size=input.element_size], ctx: DeviceContext)

Performs argsort on input buffer, storing indices in output buffer.

Parameters:

  • ascending (Bool): Sort direction (True for ascending, False for descending).
  • target (StringSlice): Target device ("cpu" or "gpu").

Args:

  • output (TileTensor): Buffer to store sorted indices.
  • input (TileTensor): Buffer containing values to sort.
  • ctx (DeviceContext): Device context for execution.

argsort[ascending: Bool = True](output: TileTensor[output.dtype, output.LayoutType, output.origin, linear_idx_type=output.linear_idx_type, element_size=output.element_size], input: TileTensor[input.dtype, input.LayoutType, input.origin, address_space=input.address_space, linear_idx_type=input.linear_idx_type, element_size=input.element_size])

CPU-only version of argsort.

Parameters:

  • ascending (Bool): Sort direction (True for ascending, False for descending).

Args:

  • output (TileTensor): Buffer to store sorted indices.
  • input (TileTensor): Buffer containing values to sort.

Was this page helpful?