Tensor and numerical foundations · Easy

Stable Vocabulary Softmax

Problem 1 of 50 · Python and NumPy

Problem

Convert each row of next-token logits into probabilities. Return a matrix of identical shape, normalizing independently across the vocabulary (the last axis).

Subtract each row maximum before exponentiating so large absolute scores remain safe. Equal scores receive equal probability.

def vocabulary_softmax(logits):

Constraints

logits is a nonempty T × V matrix; 1 ≤ T,V ≤ 32.
All values are finite and in [-10000,10000].
Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Example 1

Arguments: [
  [
    [
      0,
      0
    ]
  ]
]
Expected result: [
  [
    0.5,
    0.5
  ]
]

Both logits are zero, so each token receives 1/2.

Example 2

Arguments: [
  [
    [
      1000,
      1000,
      1000
    ]
  ]
]
Expected result: [
  [
    0.3333333333333333,
    0.3333333333333333,
    0.3333333333333333
  ]
]

Subtracting 1000 gives three zeros, so every probability is 1/3 without overflow.

Browse all coding problems · Learn the foundations