Skip to content

← Indexing and Selection step 9 of 9

Medium Primitives

Labels to one-hot

labels  [2, 0, 1]      classes 4

one-hot [[0, 0, 1, 0],
         [1, 0, 0, 0],
         [0, 1, 0, 0]]

Row i is zero everywhere except column labels[i]. This is the last indexing idea in the track and it is the one that closes the loop: gather reads at positions an index names, scatter writes at them, and one-hot is scatter with a source of ones.

The construction

out = torch.zeros(len(labels), n_classes)
out.scatter_(1, labels.unsqueeze(1), 1.0)

Three things in that line, all of which you have now met:

  • dim=1, because the class axis is the one being written along.
  • labels.unsqueeze(1) to make the index (rows, 1), the shape of what is being written. Same keepdim reasoning as everywhere else in this track.
  • a scalar source, 1.0, broadcast across every position the index names.

Note that scatter_ is right here rather than scatter_add_: within a row there is exactly one index, so no two writes collide and there is nothing to accumulate. That is the opposite of the counting problem earlier in this track, and the difference is entirely about whether the index repeats.

The one you would actually call

torch.nn.functional.one_hot(labels, num_classes)

It exists and you should use it. This problem asks for the scatter because the scatter is the thing that generalises: label smoothing, multi-hot targets, and scattering values that are not 1 all follow from it, and none of them are one_hot.

Your task

def one_hot(labels: torch.Tensor, n_classes: int) -> torch.Tensor

Return the one-hot encoding as a float tensor, in at most three dispatched operations, without calling torch.nn.functional.one_hot.

The starter loops over the labels.