Easy Primitives

Implement One-Hot Encoding

One-hot encoding represents a class ID with one row of zeros and a single one. The column containing the one is the class ID, counted from zero.

Your task

Implement one_hot(labels, num_classes) without torch.nn.functional.one_hot or jax.nn.one_hot.

labels has shape (N,). Its values are integers in 0 <= label < num_classes, delivered by the runtime as floats.
num_classes is a positive integer giving the output width, including classes that do not appear in labels.
Return a float32 tensor of shape (N, num_classes) on the input device. Leave labels unchanged.

For each input position i and class column c:

output[i, c] = 1.0 if labels[i] == c else 0.0

For labels=[2, 0, 2] and num_classes=4, return:

[[0, 0, 1, 0],
 [1, 0, 0, 0],
 [0, 0, 1, 0]]

Keep the input order and repeated labels. Each label gets its own row; do not sort, deduplicate or infer the width from the largest label.

For an empty labels vector, return shape (0, num_classes). With one class, all valid labels are zero and every output row is [1]. Negative or out-of-range labels are outside this exercise's input contract.

Two approaches

You can allocate a zero matrix and set one entry per row using integer indices. If you choose this approach, cast labels to an integer dtype before indexing. Alternatively, compare a column of labels with a row of class IDs using broadcasting, then convert the boolean result to float32.

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…