Temperature and Top-K Distribution

Problem

Return a vocabulary-wide probability vector after temperature scaling and top-k filtering. Retain exactly k tokens ranked by decreasing original logit, breaking ties by smaller token ID. Excluded tokens receive zero probability.

For temperature > 0, apply stable softmax to retained logits divided by temperature. At temperature = 0 return a one-hot vector at the smallest-ID maximum, independent of k.

def top_k_distribution(logits, temperature, k):

Constraints

• logits contains V finite values, 1 ≤ V ≤ 32; 1 ≤ k ≤ V.
• temperature is zero or in [0.001,100]; logits lie in [-10000,10000].
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    0,
    0.6931471805599453,
    1.0986122886681098
  ],
  1,
  2
]
Expected result: [
  0,
  0.4,
  0.6000000000000001
]

Arguments: [
  [
    2,
    2,
    1
  ],
  0,
  3
]
Expected result: [
  1,
  0,
  0
]