7 PyTorch Tasks (Advanced)

This capstone packs seven self-contained tasks that exercise every operation in the track. Each one is small enough to finish in a few minutes. The goal is fluency: reaching for the right operation without pausing over shapes.

The idea

The seven tasks, each a function you could keep in a utilities file:

1. Normalize a batch by subtracting the per-feature mean and dividing by the per-feature standard deviation over the batch dimension.
2. Build a causal mask that blocks attention to future tokens.
3. Split a (B, T, D) tensor into H heads, giving (B, H, T, D/H).
4. Implement scaled dot-product attention from scratch.
5. Implement a linear layer from raw tensors and match nn.Linear numerically.
6. One-hot encode integer labels of shape (B,) into (B, C).
7. Generate sinusoidal positional encodings for a sequence.

Each task needs only creation, matmul, transpose, reshape, indexing, cat/stack, and the special constructors — plus autograd once the weights become learnable. Write the expected output shape as a comment before you run anything; that habit catches most bugs before the traceback does.

Worked example

Scaled dot-product attention is the composite task. For q, k, v each (B, H, T, D):

• scores = q @ k.transpose(-2, -1) / sqrt(D) → (B, H, T, T)
• mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1) → (T, T)
• scores.masked_fill(mask, float("-inf")) → (B, H, T, T)
• attn = torch.softmax(scores, dim=-1) → (B, H, T, T)
• out = attn @ v → (B, H, T, D)
• out.transpose(1, 2).reshape(B, T, H * D) → (B, T, H * D)

The head split is the same transpose-and-reshape pair in miniature:

B, T, D, H = 2, 4, 8, 2
x = torch.randn(B, T, D)
heads = x.reshape(B, T, H, D // H).transpose(1, 2)
print(heads.shape)                 # torch.Size([2, 2, 4, 4])
merged = heads.transpose(1, 2).reshape(B, T, D)
print(torch.allclose(merged, x))   # True

In code

import math
import torch

def attention(q, k, v):
    d = q.shape[-1]
    scores = q @ k.transpose(-2, -1) / math.sqrt(d)        # (B, H, T, T)
    T = q.shape[-2]
    mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
    scores = scores.masked_fill(mask, float("-inf"))
    return torch.softmax(scores, dim=-1) @ v                # (B, H, T, D)

B, H, T, D = 2, 4, 8, 16
q, k, v = (torch.randn(B, H, T, D) for _ in range(3))
print(attention(q, k, v).shape)                             # torch.Size([2, 4, 8, 16])

Check yourself

1. In the head split, why is transpose(1, 2) needed before attention, and what goes wrong if you skip it?
2. How do you turn labels [0, 2, 1] into a (3, 3) one-hot tensor?
3. Why does the causal mask use -inf rather than 0 before the softmax?

Key takeaways

• Seven small tasks cover creation, matmul, transpose, reshape, indexing, cat/stack, and constructors.
• Write the expected shape before running; most failures are shape failures.
• Attention is just matmul plus mask plus softmax plus reshape, in that order.