Top-k truncated sampling distribution (morie.fn grtop)

Keeps the k highest probabilities (ties broken by the lower index, matching numpy's lexsort on (-p, index)), renormalises by their mass and zeroes the rest. kept_indices are 0-based and ascending.

Usage

morie_geron_topk_sampling(logits, k, T = 1)

Arguments

logits

Score vector.

k

Truncation size in [1, length(logits)].

T

Positive temperature.

Value

List with probabilities, kept_indices, kept_mass, full_probabilities, entropy.

Examples

morie_geron_topk_sampling(logits = c(1, 2, 3, 4, 5, 6, 7, 8), k = 5L)
#> $probabilities
#> [1] 0.00000000 0.00000000 0.00000000 0.01165623 0.03168492 0.08612854 0.23412166
#> [8] 0.63640865
#> 
#> $kept_indices
#> [1] 3 4 5 6 7
#> 
#> $kept_mass
#> [1] 0.9935954
#> 
#> $full_probabilities
#> [1] 0.0005766128 0.0015673960 0.0042606241 0.0115815771 0.0314819905
#> [6] 0.0855769227 0.2326221940 0.6323326828
#> 
#> $entropy
#> [1] 0.9999728
#> 
#> $estimate
#> [1] 0.00000000 0.00000000 0.00000000 0.01165623 0.03168492 0.08612854 0.23412166
#> [8] 0.63640865
#> 
#> $n
#> [1] 8
#>