Why mask before softmax, not after

Background

The right way to mask is to add -∞ to future scores before softmax — the rows still sum to 1. The tempting-but-wrong way is to softmax the full scores and then zero the future weights — now each row sums to less than 1, under-weighting the surviving values.

This problem compares the two by their row-sums.

Problem statement

Implement pre_vs_post_rowsums(scores) returning (pre_rowsums, post_rowsums).

Input

• scores — raw attention scores, shape (T, T).

Output

Returns (pre_rowsums, post_rowsums), each an np.ndarray of length T:

• pre: mask future (j > i) with -∞, row-softmax, then sum each row → all 1.0.
• post: row-softmax the full scores, then zero entries with j > i, then sum each row → generally < 1.

Examples

Example 1 — uniform scores

Input:  scores = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
Output: pre  = [1.0, 1.0, 1.0]
        post = [0.3333…, 0.6667…, 1.0]

Explanation: post-masking leaves row 0 with only 1/3 of its mass; pre-masking renormalises so every row sums to 1.

Constraints

• "Allowed" means j <= i (lower triangle incl. diagonal).
• pre: softmax over masked scores (-inf on disallowed) → rows sum to 1.
• post: softmax full scores, then zero disallowed entries → rows sum to ≤ 1.
• Return both length-T row-sum vectors.

Notes

• The exact-zero, sums-to-1 property of pre-softmax masking is why it's the correct, GPU-friendly choice — no second renormalisation needed.