Skip to main content
Log in

Mojo function

top_k_shape_impl

top_k_shape_impl[type: DType, rank: Int, single_thread_blocking_override: Bool](input: NDBuffer[type, rank, origin], k: Int, axis: Int) -> Index[rank]

Compute the output shape of a top/bottom k operation.

Parameters:

  • type (DType): Data type of the input buffer.
  • rank (Int): Rank of the input.
  • single_thread_blocking_override (Bool): If this function can block.

Args:

  • input (NDBuffer[type, rank, origin]): The input tensor.
  • k (Int): The K value in a tensor.
  • axis (Int): The axis value in a tensor.

Returns:

The output shape.