Mojo function

topk_topp_sampling_from_prob

topk_topp_sampling_from_prob[dtype: DType, out_idx_type: DType, block_size: Int = 1024, TopKArrLayoutType: TensorLayout = Layout[*?, *?], IndicesLayoutType: TensorLayout = Layout[*?, *?], TopPArrLayoutType: TensorLayout = Layout[*?, *?], SeedLayoutType: TensorLayout = Layout[*?, *?]](ctx: DeviceContext, probs: TileTensor[dtype, address_space=probs.address_space, linear_idx_type=probs.linear_idx_type, element_size=probs.element_size], output: TileTensor[out_idx_type, address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size], top_k_val: Int, top_p_val: Float32 = 1, deterministic: Bool = False, rng_seed: Optional[TileTensor[DType.uint64, SeedLayoutType, ImmutAnyOrigin]] = None, rng_offset: UInt64 = UInt64(0), indices: Optional[TileTensor[out_idx_type, IndicesLayoutType, ImmutAnyOrigin]] = None, top_k_arr: Optional[TileTensor[out_idx_type, TopKArrLayoutType, ImmutAnyOrigin]] = None, top_p_arr: Optional[TileTensor[DType.float32, TopPArrLayoutType, ImmutAnyOrigin]] = None)

Joint top-k + top-p sampling from probability distribution.

Performs stochastic sampling considering only tokens that satisfy both the top-k count constraint AND the top-p nucleus constraint. When top_p_val is 1.0 (default) this behaves identically to topk_sampling_from_prob.

Args:

ctx (DeviceContext): Device context for kernel execution.
probs (TileTensor[dtype, address_space=probs.address_space, linear_idx_type=probs.linear_idx_type, element_size=probs.element_size]): Input probability distribution [batch_size, d].
output (TileTensor[out_idx_type, address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size]): Output sampled indices [batch_size].
top_k_val (Int): Default top-k value (number of top tokens to consider).
top_p_val (Float32): Default top-p value (nucleus probability threshold).
deterministic (Bool): Whether to use deterministic sampling.
rng_seed (Optional[TileTensor[DType.uint64, SeedLayoutType, ImmutAnyOrigin]]): Optional seed tensor. If provided, rng_seed[0] is used as the seed. If None, defaults to 0.
rng_offset (UInt64): Random offset for Random number generator.
indices (Optional[TileTensor[out_idx_type, IndicesLayoutType, ImmutAnyOrigin]]): Optional row indices for batch indexing [batch_size].
top_k_arr (Optional[TileTensor[out_idx_type, TopKArrLayoutType, ImmutAnyOrigin]]): Optional per-row top-k values [batch_size].
top_p_arr (Optional[TileTensor[DType.float32, TopPArrLayoutType, ImmutAnyOrigin]]): Optional per-row top-p values [batch_size].

Raises:

Error: If tensor ranks or shapes are invalid.