Mask the padding before the softmax

A Boolean mask can identify attention-score positions to exclude. In this exercise, True means replace the score and False means keep it:

import torch

scores = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
mask = torch.tensor([[False, False, True], [False, True, True]])
out = scores.masked_fill(mask, float("-inf"))
print(out.tolist())  # [[1.0, 2.0, -inf], [4.0, -inf, -inf]]
print(scores.tolist())  # [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]

For finite unmasked scores and at least one valid position per row, negative infinity gives excluded positions zero probability under softmax. The task here only fills scores; it does not compute softmax.

Apply the mask as a tensor operation

The starter loops over every score and converts each mask entry to a Python Boolean. That introduces repeated indexing and dispatch overhead, and reading device values into Python can require synchronization. masked_fill expresses the whole replacement directly.

torch.where(mask, replacement, scores) is another option, particularly when replacements vary by position. A scalar replacement works here too. Both forms can broadcast compatible shapes, though this exercise supplies scores and mask with exactly matching shapes.

The harness counts dispatched PyTorch operations, not GPU kernels. Both forms fit its two-operation limit on the tested inputs; one Python API call does not necessarily mean one dispatched operation.

Preserve the input

masked_fill returns a new tensor. masked_fill_ changes the original, so it does not satisfy this task. In-place mutation is appropriate only when the surrounding code permits it, including any autograd constraints.

Negative infinity and fully masked rows

A finite value such as -1e9 is not interchangeable with negative infinity. It is outside float16's finite range, and passing it as a fill value for a float16 tensor can raise an overflow error. A finite masking value also need not exclude a position relative to every possible unmasked score.

If all scores in a row become negative infinity, a plain softmax returns NaNs. Filling all positions with the same finite negative value instead gives a uniform distribution, which also does not mean every position was excluded. Define a policy for fully masked rows before using softmax. The all-True test here expects a row of negative infinities, without softmax.

Select values using coordinates

A slice selects a regularly spaced region of an array. A gather selects entries using supplied indices. Check whether those indices select whole rows or columns, or a different entry within each row.

For [[10, 11], [20, 21]], choosing column 1 from the first row and column 0 from the second gives [11, 20]. Choosing columns [1, 0] from every row instead gives [[11, 10], [21, 20]]. These operations use the same index values but have different meanings and output shapes.

Scatter writes values to selected destinations

Scatter takes values and destination indices. The operation may replace existing values or combine them using a reduction such as addition.

Starting with [0, 0, 0], scatter-add of values [2, 3] at indices [1, 1] gives [0, 5, 0]. Both values contribute to the repeated destination. Repeated assignment is not equivalent to this sum; follow the operation's rule for repeated indices.

Selecting with a mask differs from replacing with a mask

For values [10, 20, 30] and mask [True, False, True], selecting masked entries gives [10, 30]. Replacing masked entries with zero gives [0, 20, 0]. Selection can change the shape; replacement preserves it. A selection whose size depends on runtime values needs particular care in compiled code with fixed output shapes.

Check whether an update changes the original array

JAX indexed updates return a new array value. In PyTorch, basic slicing can return a view that shares storage, while advanced indexing returns a copy. Modifying a selected copy does not modify the source. Direct indexed assignment is a separate operation, so distinguish changing a selection from assigning into the original.