Mojo function

topk_sampling_from_prob

topk_sampling_from_prob[dtype: DType, out_idx_type: DType, block_size: Int = 1024, top_k_arr_shape_types: Variadic[CoordLike] = RuntimeInt[DType.int64], top_k_arr_stride_types: Variadic[CoordLike] = ComptimeInt[1], indices_shape_types: Variadic[CoordLike] = RuntimeInt[DType.int64], indices_stride_types: Variadic[CoordLike] = ComptimeInt[1]](ctx: DeviceContext, probs: TileTensor[dtype, origin, address_space=address_space, linear_idx_type=linear_idx_type, element_shape_types=element_shape_types], output: TileTensor[out_idx_type, origin, address_space=address_space, linear_idx_type=linear_idx_type, element_shape_types=element_shape_types], top_k_val: Int, deterministic: Bool = False, rng_seed: UInt64 = 0, rng_offset: UInt64 = 0, indices: Optional[TileTensor[out_idx_type, MutExternalOrigin]] = None, top_k_arr: Optional[TileTensor[out_idx_type, MutExternalOrigin]] = None)

Top-K sampling from probability distribution.

Performs stochastic sampling from a probability distribution, considering only the top-k most probable tokens. Uses rejection sampling with ternary search to efficiently find appropriate samples.

Args:

ctx (DeviceContext): Device context for kernel execution.
probs (TileTensor): Input probability distribution [batch_size, d].
output (TileTensor): Output sampled indices [batch_size].
top_k_val (Int): Default top-k value (number of top tokens to consider).
deterministic (Bool): Whether to use deterministic sampling.
rng_seed (UInt64): Random seed for Random number generator.
rng_offset (UInt64): Random offset for Random number generator.
indices (Optional): Optional row indices for batch indexing [batch_size].
top_k_arr (Optional): Optional per-row top-k values [batch_size].

Raises:

Error: If tensor ranks or shapes are invalid.