Medium Framework

Masked Fill

A mask selects which tensor entries to replace. Here, true means replace with a supplied scalar; false means keep the original value.

Your task

Implement masked_fill(x, mask, fill_value).

x is a floating-point tensor of any shape.
mask has exactly the same shape. Test inputs contain 0.0 and 1.0; convert these to boolean values before using them as a condition.
fill_value is a scalar replacement value, not a tensor to broadcast.

Return a tensor with the same shape and dtype as x, leaving both inputs unchanged. Apply this rule independently to each entry:

output[index] = fill_value if mask[index] is true else x[index]

For x=[[1, 2], [3, 4]], mask=[[0, 1], [1, 0]] and fill_value=-1, return [[1, -1], [-1, 4]]. Keep all positions in the output; do not gather only the selected entries.

An all-false mask keeps every value. An all-true mask replaces every value. For an empty input, return an empty tensor of the same shape.

Applying the mask

torch.where(condition, fill_value, x) and jnp.where(condition, fill_value, x) select the replacement when the condition is true. PyTorch's nonmutating x.masked_fill(condition, fill_value) also implements this rule. Pass the converted boolean mask as condition.

Multiplying x by a mask is a different operation: it cannot insert an arbitrary replacement value. Also check the mask convention when using this operation elsewhere; some APIs use true to mean keep rather than replace.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

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.

Detecting runtime…