Frequency tiers (estimates, relative to interviews targeting this pack's specialty — a ★★★★★ here means "core for these roles", not for generalist ML loops): ★★★★★ expect it in most loops · ★★★★ very common · ★★★ common · ★★ occasional · ★ rare/deep-specialist. Spend ~50% of practice on the top two tiers.
70+ problems across foundations, attention, optimizers, vision, RL, diffusion, PEFT Each with full PyTorch / NumPy solution and test cases
Principal/Senior-Principal CV/ML Interview Prep
Notes
This pack covers the questions that come up in principal-level ML / CV interviews and in research-engineer hiring at FAANG, top labs, and AI start-ups. Each problem has: (i) a precise statement, (ii) a clean reference solution in NumPy or PyTorch, (iii) test cases that verify correctness against either the framework’s reference op or known closed-form behaviour. Read the solutions side by side with PyTorch’s documentation: the goal is not to memorise but to be able to derive these in under 10 minutes on a whiteboard.
52 DDPM reverse step (µ, σ2 closed form)
Problem
Implement softmax in NumPy along an arbitrary axis without overflow for large logits. Matchtorch.softmax to within 1×10−6.
import numpy as np
def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
x_max = np.max(x, axis=axis, keepdims=True)
e = np.exp(x - x_max) # subtract max for stability
return e / np.sum(e, axis=axis, keepdims=True)
Tests
import torch
x = np.array([[1000., 1001., 1002.], [-50., 0., 50.]])
out = softmax(x, axis=-1)
ref = torch.softmax(torch.tensor(x), dim=-1).numpy()
assert np.allclose(out, ref, atol=1e-6)
assert np.allclose(out.sum(axis=-1), 1.0)
print("softmax OK")
Problem
Implement logsumexp(x, axis) stably; this is the building block for log-softmax and cross-entropy.
import numpy as np
def logsumexp(x: np.ndarray, axis: int = -1) -> np.ndarray:
m = np.max(x, axis=axis, keepdims=True)
return (np.log(np.sum(np.exp(x - m), axis=axis, keepdims=True)) + m).squeeze(axis=axis)
Tests
from scipy.special import logsumexp as scipy_lse
x = np.random.randn(4, 5) * 1000.
assert np.allclose(logsumexp(x, axis=1), scipy_lse(x, axis=1))
print("logsumexp OK")
Problem
Implement log_softmax(x, axis) without computing softmax then taking log.
import numpy as np
def log_softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
m = np.max(x, axis=axis, keepdims=True)
z = x - m
lse = np.log(np.sum(np.exp(z), axis=axis, keepdims=True))
return z - lse
Tests
x = np.random.randn(3, 7) * 50.
ref = torch.log_softmax(torch.tensor(x), dim=-1).numpy()
assert np.allclose(log_softmax(x, -1), ref, atol=1e-6)
print("log_softmax OK")
Problem
Implement sigmoid(x) that does not overflow for large positive or negative x.
import numpy as np
def sigmoid(x: np.ndarray) -> np.ndarray:
out = np.empty_like(x, dtype=np.float64)
pos = x >= 0
out[pos] = 1.0 / (1.0 + np.exp(-x[pos]))
e = np.exp(x[~pos])
out[~pos] = e / (1.0 + e)
return out
Tests
x = np.array([-1000., -10., 0., 10., 1000.])
out = sigmoid(x)
ref = torch.sigmoid(torch.tensor(x)).numpy()
assert np.allclose(out, ref, atol=1e-12)
print("sigmoid OK")
Problem
Given logits z of shape (N, C) and integer targets y of shape (N, ), compute the mean cross-entropy. Match torch.nn.functional.cross_entropy.
import numpy as np
def cross_entropy(logits: np.ndarray, targets: np.ndarray) -> float:
log_p = log_softmax(logits, axis=-1)
n = logits.shape[0]
return -log_p[np.arange(n), targets].mean()
Tests
import torch.nn.functional as F
N, C = 8, 5
z = np.random.randn(N, C) * 5.
y = np.random.randint(0, C, size=(N,))
out = cross_entropy(z, y)
ref = F.cross_entropy(torch.tensor(z), torch.tensor(y)).item()
assert abs(out - ref) < 1e-6
print("cross_entropy OK")
Problem
Implement BCE-with-logits using log(1 + ex) = softplus(x) and the identity log σ(x) = −softplus(−x).
import numpy as np
def softplus(x):
# numerically stable softplus
return np.where(x > 0, x + np.log1p(np.exp(-x)), np.log1p(np.exp(x)))
def bce_with_logits(z: np.ndarray, y: np.ndarray) -> float:
# -[y*log(sigmoid(z)) + (1-y)*log(1-sigmoid(z))]
return (softplus(z) - y * z).mean()
Tests
z = np.random.randn(20) * 10
y = (np.random.rand(20) > 0.5).astype(np.float64)
ref = F.binary_cross_entropy_with_logits(torch.tensor(z), torch.tensor(y)).item()
assert abs(bce_with_logits(z, y) - ref) < 1e-7
print("bce OK")
Problem
Implement relu and leaky_relu(slope=0.01) in NumPy. Provide their derivatives.
import numpy as np
def relu(x): return np.maximum(0, x)
def relu_grad(x): return (x > 0).astype(x.dtype)
def leaky_relu(x, s=0.01): return np.where(x >= 0, x, s * x)
def leaky_relu_grad(x, s=0.01): return np.where(x >= 0, 1.0, s)
Tests
x = np.array([-1.0, 0.0, 1.0])
assert np.allclose(relu(x), [0,0,1]) and np.allclose(leaky_relu(x), [-0.01,0,1])
print("relu/leaky OK")
Problem
x · 1√Implement exact GELU xΦ(x) and the tanh approximation used by GPT-2:2(1 + tanh(2/π(x + 0.044715x3))).
import math
def gelu_exact(x):
return 0.5 * x * (1.0 + np.erf(x / math.sqrt(2.0))) # uses np.erf via scipy or numpy>=2
def gelu_tanh(x):
c = math.sqrt(2.0 / math.pi)
return 0.5 * x * (1.0 + np.tanh(c * (x + 0.044715 * x ** 3)))
Tests
from scipy.special import erf
def gelu_exact(x): return 0.5 * x * (1.0 + erf(x / math.sqrt(2.0)))
x = np.linspace(-3, 3, 13)
ref = F.gelu(torch.tensor(x), approximate='none').numpy()
assert np.allclose(gelu_exact(x), ref, atol=1e-7)
ref2 = F.gelu(torch.tensor(x), approximate='tanh').numpy()
assert np.allclose(gelu_tanh(x), ref2, atol=1e-6)
print("gelu OK")
Problem
Implement silu(x) = xσ(x) (also called Swish).
def silu(x): return x * sigmoid(x)
Tests
x = np.linspace(-5, 5, 11)
assert np.allclose(silu(x), F.silu(torch.tensor(x)).numpy(), atol=1e-7)
print("silu OK")
Problem
Implement the SwiGLU FFN used in LLaMA / Mistral: SwiGLU(x) = (silu(xW1) ⊙xW2) W3.
import torch.nn as nn
class SwiGLU(nn.Module):
def __init__(self, dim, hidden):
super().__init__()
self.w1 = nn.Linear(dim, hidden, bias=False) # gate
self.w2 = nn.Linear(dim, hidden, bias=False) # up
self.w3 = nn.Linear(hidden, dim, bias=False) # down
def forward(self, x):
return self.w3(F.silu(self.w1(x)) * self.w2(x))
Tests
m = SwiGLU(64, 128)
y = m(torch.randn(2, 10, 64))
assert y.shape == (2, 10, 64)
print("SwiGLU OK")
Problem
Implement Linear with explicit forward and backward (computing gradients w.r.t. input, weight, and bias).
import numpy as np
class LinearManual:
def __init__(self, in_f, out_f):
self.W = np.random.randn(out_f, in_f) * np.sqrt(2.0 / in_f)
self.b = np.zeros(out_f)
def forward(self, x):
self.x = x # (B, in_f)
return x @ self.W.T + self.b # (B, out_f)
def backward(self, dy): # dy: (B, out_f)
dx = dy @ self.W
dW = dy.T @ self.x # (out_f, in_f)
db = dy.sum(axis=0)
return dx, dW, db
Tests
B, I, O = 4, 5, 3
m = LinearManual(I, O)
x = np.random.randn(B, I)
y = m.forward(x)
dy = np.random.randn(B, O)
dx, dW, db = m.backward(dy)
# numerical check on dx via torch autograd
t_x = torch.tensor(x, requires_grad=True)
t_W = torch.tensor(m.W, requires_grad=True)
t_b = torch.tensor(m.b, requires_grad=True)
t_y = t_x @ t_W.T + t_b
t_y.backward(torch.tensor(dy))
assert np.allclose(dx, t_x.grad.numpy(), atol=1e-7)
assert np.allclose(dW, t_W.grad.numpy(), atol=1e-7)
assert np.allclose(db, t_b.grad.numpy(), atol=1e-7)
print("Linear OK")
Problem
Implement a 2D convolution forward via the im2col trick. Stride, padding, single batch dimension.
import numpy as np
def im2col(x, kh, kw, stride=1, pad=0):
B, C, H, W = x.shape
Hp, Wp = H + 2*pad, W + 2*pad
xp = np.pad(x, ((0,0),(0,0),(pad,pad),(pad,pad)))
out_h = (Hp - kh) // stride + 1
out_w = (Wp - kw) // stride + 1
cols = np.zeros((B, C, kh, kw, out_h, out_w))
for i in range(kh):
for j in range(kw):
cols[:, :, i, j, :, :] = xp[:, :, i:i+stride*out_h:stride, j:j+stride*out_w:stride]
return cols.transpose(0, 4, 5, 1, 2, 3).reshape(B*out_h*out_w, -1), out_h, out_w
def conv2d_forward(x, W, b, stride=1, pad=0):
# W: (Cout, Cin, kh, kw)
Cout, Cin, kh, kw = W.shape
cols, oh, ow = im2col(x, kh, kw, stride, pad)
out = cols @ W.reshape(Cout, -1).T + b
return out.reshape(x.shape[0], oh, ow, Cout).transpose(0, 3, 1, 2)
Tests
x = np.random.randn(2, 3, 8, 8)
W = np.random.randn(4, 3, 3, 3); b = np.random.randn(4)
out = conv2d_forward(x, W, b, stride=1, pad=1)
ref = F.conv2d(torch.tensor(x), torch.tensor(W), torch.tensor(b), stride=1, padding=1).numpy()
assert np.allclose(out, ref, atol=1e-6)
print("conv2d OK")
Advanced implementation. sliding_window_view extracts all patches as a free stride-trick view (no copy) and the whole convolution collapses to one einsum — ~5x faster than im2col and zero Python loops. Verified equal to F.conv2d (1e-6).
import numpy as np
from numpy.lib.stride_tricks import sliding_window_view
def conv2d_einsum(x, w, b, stride=1, pad=0):
xp = np.pad(x, ((0, 0), (0, 0), (pad, pad), (pad, pad)))
win = sliding_window_view(xp, (w.shape[2], w.shape[3]), axis=(2, 3))
win = win[:, :, ::stride, ::stride] # (N, C, Ho, Wo, kH, kW), still a view
return np.einsum("nchwkl,fckl->nfhw", win, w, optimize=True) + b[None, :, None, None]
Problem
Implement 2D max pooling forward, recording argmax for the backward pass.
import numpy as np
def maxpool2d_forward(x, k=2, s=2):
B, C, H, W = x.shape
out_h, out_w = H // s, W // s
out = np.zeros((B, C, out_h, out_w))
arg = np.zeros_like(out, dtype=np.int64)
for i in range(out_h):
for j in range(out_w):
patch = x[:, :, i*s:i*s+k, j*s:j*s+k].reshape(B, C, -1)
out[:, :, i, j] = patch.max(axis=-1)
arg[:, :, i, j] = patch.argmax(axis=-1)
return out, arg
Tests
x = np.random.randn(1, 1, 4, 4)
out, _ = maxpool2d_forward(x, 2, 2)
ref = F.max_pool2d(torch.tensor(x), 2, 2).numpy()
assert np.allclose(out, ref)
print("maxpool OK")
Problem
Implement inverted dropout: during training scale by 1/(1 −p), during eval pass through.
import numpy as np
def dropout(x, p=0.1, train=True, rng=None):
if (not train) or p == 0.0:
return x, None
rng = rng or np.random
mask = (rng.rand(*x.shape) >= p).astype(x.dtype) / (1.0 - p)
return x * mask, mask
Tests
np.random.seed(0)
x = np.ones((1000,))
y, _ = dropout(x, p=0.3, train=True)
assert abs(y.mean() - 1.0) < 0.05 # expectation preserved
y2, _ = dropout(x, p=0.3, train=False)
assert np.allclose(y2, x) # eval = identity
print("dropout OK")
Problem
Implement training-mode BatchNorm1d forward: per-feature mean/var across the batch, with learnable affine γ, β.
import numpy as np
def batchnorm_forward(x, gamma, beta, eps=1e-5):
mu = x.mean(axis=0)
var = x.var(axis=0)
xhat = (x - mu) / np.sqrt(var + eps)
return gamma * xhat + beta, (mu, var, xhat)
Tests
B, F_ = 16, 8
x = np.random.randn(B, F_).astype(np.float32)
gamma = np.ones(F_, dtype=np.float32); beta = np.zeros(F_, dtype=np.float32)
y, _ = batchnorm_forward(x, gamma, beta)
m = nn.BatchNorm1d(F_, affine=False).eval() # use functional w/o running stats
ref = F.batch_norm(torch.tensor(x), None, None, training=True)
assert np.allclose(y, ref.numpy(), atol=1e-5)
print("BN OK")
Problem
Implement LayerNorm normalising over the last axis, with affine γ, β.
import numpy as np
def layernorm(x, gamma, beta, eps=1e-5):
mu = x.mean(axis=-1, keepdims=True)
var = x.var(axis=-1, keepdims=True)
return gamma * (x - mu) / np.sqrt(var + eps) + beta
Tests
B, T, D = 2, 5, 8
x = np.random.randn(B, T, D).astype(np.float32)
g = np.ones(D, dtype=np.float32); b = np.zeros(D, dtype=np.float32)
out = layernorm(x, g, b)
ref = F.layer_norm(torch.tensor(x), (D,)).numpy()
assert np.allclose(out, ref, atol=1e-5)
print("LN OK")
Problem
RMSNorm (Zhang & Sennrich 2019, used in LLaMA/Mistral): no mean subtraction, divide by RMS.
x RMSNorm(x) = γ ·. q 1i x2Σi + ε D
import numpy as np
import torch
import torch.nn as nn
def rmsnorm(x, gamma, eps=1e-6):
rms = np.sqrt((x ** 2).mean(axis=-1, keepdims=True) + eps)
return gamma * x / rms
class RMSNorm(nn.Module):
def __init__(self, d, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(d))
self.eps = eps
def forward(self, x):
rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
return self.weight * x / rms
Tests
x = torch.randn(3, 4, 16)
m = RMSNorm(16); y = m(x)
assert torch.allclose(y.std(dim=-1, unbiased=False).clamp_min(0), torch.ones(3,4), atol=1e-2)
print("RMSNorm OK")
Problem
GroupNorm splits channels into G groups; normalises within each group along (group, H, W).
import numpy as np
def groupnorm(x, G, gamma, beta, eps=1e-5):
B, C, H, W = x.shape
x = x.reshape(B, G, C // G, H, W)
mu = x.mean(axis=(2,3,4), keepdims=True)
var = x.var(axis=(2,3,4), keepdims=True)
xhat = (x - mu) / np.sqrt(var + eps)
xhat = xhat.reshape(B, C, H, W)
return gamma.reshape(1, C, 1, 1) * xhat + beta.reshape(1, C, 1, 1)
Tests
B, C, H, W = 2, 8, 4, 4; G = 4
x = np.random.randn(B, C, H, W).astype(np.float32)
g = np.ones(C, dtype=np.float32); b = np.zeros(C, dtype=np.float32)
out = groupnorm(x, G, g, b)
ref = F.group_norm(torch.tensor(x), G).numpy()
assert np.allclose(out, ref, atol=1e-5)
print("GN OK")
Problem
Implement plain SGD, classic momentum, and Nesterov.
import numpy as np
class SGD:
def __init__(self, params, lr=1e-2, mom=0.0, nesterov=False, wd=0.0):
self.params = list(params)
self.lr, self.mom, self.nest, self.wd = lr, mom, nesterov, wd
self.v = [np.zeros_like(p) for p in self.params]
def step(self, grads):
for i, (p, g) in enumerate(zip(self.params, grads)):
if self.wd: g = g + self.wd * p
self.v[i] = self.mom * self.v[i] + g
update = (self.mom * self.v[i] + g) if self.nest else self.v[i]
p -= self.lr * update
Tests
p = np.array([1.0, 2.0])
g = np.array([0.1, -0.2])
opt = SGD([p], lr=0.5, mom=0.9)
opt.step([g]); opt.step([g])
# v_2 = 0.9*v_1 + g, v_1 = g; expected change visible
print("sgd step OK", p)
Problem
Implement Adam (Kingma & Ba 2015): bias-corrected first/second moments + element-wise LR.
import numpy as np
class Adam:
def __init__(self, params, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
self.params = list(params)
self.lr, self.b1, self.b2, self.eps = lr, b1, b2, eps
self.m = [np.zeros_like(p) for p in self.params]
self.v = [np.zeros_like(p) for p in self.params]
self.t = 0
def step(self, grads):
self.t += 1
for i, (p, g) in enumerate(zip(self.params, grads)):
self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g
self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (g * g)
mhat = self.m[i] / (1 - self.b1 ** self.t)
vhat = self.v[i] / (1 - self.b2 ** self.t)
p -= self.lr * mhat / (np.sqrt(vhat) + self.eps)
Tests
p = np.zeros(3); g = np.array([1.0, -1.0, 0.5])
opt = Adam([p], lr=1e-1)
for _ in range(50): opt.step([g])
assert np.sign(-p[0]) == np.sign(g[0]); print("Adam OK", p)
Problem
Implement AdamW (Loshchilov & Hutter 2019): weight decay applied directly to parameters, decoupled from the gradient.
import numpy as np
class AdamW(Adam):
def __init__(self, params, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, wd=1e-2):
super().__init__(params, lr, b1, b2, eps)
self.wd = wd
def step(self, grads):
self.t += 1
for i, (p, g) in enumerate(zip(self.params, grads)):
self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g
self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (g * g)
mhat = self.m[i] / (1 - self.b1 ** self.t)
vhat = self.v[i] / (1 - self.b2 ** self.t)
p -= self.lr * (mhat / (np.sqrt(vhat) + self.eps) + self.wd * p)
Tests
p = np.array([10.0]); opt = AdamW([p], lr=1e-1, wd=1e-2)
for _ in range(20): opt.step([np.zeros_like(p)])
assert p[0] < 10.0 # decoupled weight decay shrinks parameters even with zero grad
print("AdamW OK")
Advanced implementation. The per-parameter Python loop becomes one horizontally-fused multi-tensor call per operation — exactly what torch.optim's foreach=True (default) and fused=True modes do. Verified: trajectory matches torch.optim.AdamW to 1e-5.
import torch
@torch.no_grad()
def adamw_foreach_step(params, grads, m, v, t, lr, b1, b2, eps, wd):
torch._foreach_mul_(params, 1 - lr * wd) # decoupled decay first (torch order)
torch._foreach_mul_(m, b1); torch._foreach_add_(m, grads, alpha=1 - b1)
torch._foreach_mul_(v, b2); torch._foreach_addcmul_(v, grads, grads, value=1 - b2)
mhat = torch._foreach_div(m, 1 - b1 ** t)
vhat = torch._foreach_div(v, 1 - b2 ** t)
denom = torch._foreach_add(torch._foreach_sqrt(vhat), eps)
torch._foreach_addcdiv_(params, mhat, denom, value=-lr)
Problem
Implement learning-rate schedule used by every modern Transformer: linear warmup for Tw steps, then half- cycle cosine decay to a min LR.
import math
def lr_cosine(step, lr_max, warmup, total, lr_min=0.0):
if step < warmup:
return lr_max * step / max(1, warmup)
p = (step - warmup) / max(1, total - warmup)
return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * p))
Tests
assert lr_cosine(0, 1.0, 10, 100) == 0.0
assert lr_cosine(10, 1.0, 10, 100) == 1.0
assert abs(lr_cosine(100, 1.0, 10, 100, lr_min=0.1) - 0.1) < 1e-9
print("cosine OK")
Problem
Implement scaled dot-product attention with optional mask. Match F.scaled_dot_product_attention.
import math
import torch
def sdp_attention(Q, K, V, mask=None):
# Q,K,V: (B, H, T, d_k). mask: (B, 1, T, T) with 0 (keep) / -inf (mask)
d_k = Q.shape[-1]
scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)
if mask is not None: scores = scores + mask
attn = torch.softmax(scores, dim=-1)
return attn @ V, attn
Tests
B, H, T, d = 2, 4, 8, 16
Q = torch.randn(B, H, T, d); K = torch.randn(B, H, T, d); V = torch.randn(B, H, T, d)
ours, _ = sdp_attention(Q, K, V)
ref = F.scaled_dot_product_attention(Q, K, V)
assert torch.allclose(ours, ref, atol=1e-5)
print("SDPA OK")
Problem
Implement annn.Module that does multi-head attention: project Q/K/V via separate linears, split into heads, attend, concat, output project.
import torch.nn as nn
class MHA(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
assert d_model % n_heads == 0
self.h, self.dk = n_heads, d_model // n_heads
self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
self.proj = nn.Linear(d_model, d_model, bias=False)
def forward(self, x, mask=None):
B, T, D = x.shape
q, k, v = self.qkv(x).chunk(3, dim=-1) # each (B, T, D)
q = q.view(B, T, self.h, self.dk).transpose(1, 2) # (B, h, T, dk)
k = k.view(B, T, self.h, self.dk).transpose(1, 2)
v = v.view(B, T, self.h, self.dk).transpose(1, 2)
out, _ = sdp_attention(q, k, v, mask)
out = out.transpose(1, 2).contiguous().view(B, T, D)
return self.proj(out)
Tests
m = MHA(64, 4); x = torch.randn(2, 7, 64)
y = m(x); assert y.shape == x.shape; print("MHA OK")
Advanced implementation. Two einsum contractions whose index strings are the shape documentation; generalizes to GQA with one extra group index, and biases just add to att. In production, call F.scaled_dot_product_attention (FlashAttention kernels, no (T,T) matrix). Verified equal to SDPA for square and cached single-query shapes (1e-6).
import math, torch
def attention_einsum(q, k, v, causal=True): # (B, h, T, d) -> (B, h, T, d)
att = torch.einsum("bhtd,bhsd->bhts", q, k) / math.sqrt(q.size(-1))
if causal:
T, S = q.size(2), k.size(2)
mask = torch.triu(torch.ones(T, S, dtype=torch.bool, device=q.device),
diagonal=1 + S - T)
att = att.masked_fill(mask, float("-inf"))
return torch.einsum("bhts,bhsd->bhtd", att.softmax(-1), v)
Problem
Generate an additive causal mask of shape (T, T) with 0 on the lower triangle and −∞above the diagonal.
import torch
def causal_mask(T, device='cpu'):
return torch.triu(torch.full((T, T), float('-inf'), device=device), diagonal=1)
Tests
M = causal_mask(4)
assert M[0, 1] == float('-inf') and M[1, 0] == 0.0
print("causal mask OK")
Problem
Implement a single decoding step that appends the new token’s K/V to an existing cache and runs attention against the full cached sequence.
import torch
import torch.nn as nn
class CachedAttention(nn.Module):
def __init__(self, d, h):
super().__init__()
self.h, self.dk = h, d // h
self.qkv = nn.Linear(d, 3 * d, bias=False); self.proj = nn.Linear(d, d, bias=False)
def step(self, x, kv_cache):
# x: (B, 1, D) one new token
B, _, D = x.shape
q, k, v = self.qkv(x).chunk(3, dim=-1)
q = q.view(B, 1, self.h, self.dk).transpose(1, 2)
k = k.view(B, 1, self.h, self.dk).transpose(1, 2)
v = v.view(B, 1, self.h, self.dk).transpose(1, 2)
if kv_cache is None:
K, V = k, v
else:
K = torch.cat([kv_cache[0], k], dim=2) # (B, h, T+1, dk)
V = torch.cat([kv_cache[1], v], dim=2)
out, _ = sdp_attention(q, K, V) # no mask: only attending to past
out = out.transpose(1, 2).contiguous().view(B, 1, D)
return self.proj(out), (K, V)
Tests
m = CachedAttention(64, 4); cache = None
for t in range(5):
y, cache = m.step(torch.randn(1, 1, 64), cache)
assert cache[0].shape == (1, 4, 5, 16); print("KV cache OK")
Advanced implementation. torch.cat per decode step reallocates and copies the whole prefix — O(T²) total memory traffic. Preallocate max_len once and write in place; attention reads a zero-copy slice. (PagedAttention goes further: fixed-size blocks + a page table.) Verified equal to the cat-based cache at every step.
import torch
class PreallocKV:
def __init__(self, B, h, max_len, dk):
self.k = torch.zeros(B, h, max_len, dk)
self.v = torch.zeros(B, h, max_len, dk)
self.len = 0
def append(self, k_new, v_new):
t = k_new.size(2)
self.k[:, :, self.len:self.len + t] = k_new # in-place write
self.v[:, :, self.len:self.len + t] = v_new
self.len += t
return self.k[:, :, :self.len], self.v[:, :, :self.len] # views, no copy
Problem
Apply RoPE: rotate pairs of features by an angle θi · pos before attention.
import torch
def rope_freqs(dim, max_pos, base=10000.0):
freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
pos = torch.arange(max_pos).float()
angles = torch.outer(pos, freqs) # (max_pos, dim/2)
return torch.cos(angles), torch.sin(angles)
def apply_rope(x, cos, sin):
# x: (..., T, D)
x1, x2 = x[..., 0::2], x[..., 1::2]
cos = cos.to(x.device); sin = sin.to(x.device)
x_rot = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
return x_rot.flatten(-2)
Tests
T, D = 4, 8
cos, sin = rope_freqs(D, T)
x = torch.randn(1, T, D); y = apply_rope(x, cos, sin)
assert y.shape == x.shape
# Property: same content at different positions yields different rotated vectors.
y0 = apply_rope(x[:, :1], cos[:1], sin[:1])
y1 = apply_rope(x[:, :1], cos[1:2], sin[1:2])
assert not torch.allclose(y0, y1)
print("RoPE OK")
Problem
Implement ALiBi (Press et al. 2022): add a per-head linear bias −mh · |i −j| to attention scores instead of positional embeddings.
import math
import torch
def alibi_slopes(n_heads):
# geometric slope schedule used in the paper
start = 2 ** (-2 ** -(math.log2(n_heads) - 3))
return torch.tensor([start * start ** i for i in range(n_heads)])
def alibi_bias(T, n_heads):
slopes = alibi_slopes(n_heads) # (h,)
pos = torch.arange(T)
rel = pos.unsqueeze(0) - pos.unsqueeze(1) # (T, T) j - i
return -slopes.view(n_heads, 1, 1) * rel.abs().float()
Tests
B, h, T = 1, 8, 6
bias = alibi_bias(T, h).unsqueeze(0)
Q = K = V = torch.randn(B, h, T, 16)
out, _ = sdp_attention(Q, K, V, mask=bias)
assert out.shape == Q.shape; print("ALiBi OK")
Problem
Implement GQA with H query heads, G KV groups (G|H). Each KV group is shared by H/G query heads. Used by LLaMA-2/3, Mistral, Gemma.
import torch.nn as nn
class GQA(nn.Module):
def __init__(self, d, h_q, h_kv):
super().__init__()
assert d % h_q == 0 and h_q % h_kv == 0
self.h_q, self.h_kv = h_q, h_kv
self.dk = d // h_q
self.q = nn.Linear(d, h_q * self.dk, bias=False)
self.k = nn.Linear(d, h_kv * self.dk, bias=False)
self.v = nn.Linear(d, h_kv * self.dk, bias=False)
self.proj = nn.Linear(d, d, bias=False)
def forward(self, x):
B, T, D = x.shape
q = self.q(x).view(B, T, self.h_q, self.dk).transpose(1, 2)
k = self.k(x).view(B, T, self.h_kv, self.dk).transpose(1, 2)
v = self.v(x).view(B, T, self.h_kv, self.dk).transpose(1, 2)
# repeat K/V to match query head count
rep = self.h_q // self.h_kv
k = k.repeat_interleave(rep, dim=1)
v = v.repeat_interleave(rep, dim=1)
out, _ = sdp_attention(q, k, v)
return self.proj(out.transpose(1, 2).contiguous().view(B, T, D))
Tests
m = GQA(64, h_q=8, h_kv=2); x = torch.randn(2, 5, 64)
assert m(x).shape == x.shape; print("GQA OK")
Problem
Show the online-softmax trick that FlashAttention uses: stream K, V in tiles, maintain running (m, l, O) state per query.
import math
import torch
def flash_attn_naive(Q, K, V, tile=64):
# Q,K,V: (T, d). Single-head, illustrative (not fast).
Tq, d = Q.shape
Tk, _ = K.shape
O = torch.zeros(Tq, d)
m = torch.full((Tq,), -float('inf'))
l = torch.zeros(Tq)
for j in range(0, Tk, tile):
Kj = K[j:j+tile]; Vj = V[j:j+tile]
S = (Q @ Kj.T) / math.sqrt(d) # (Tq, tile)
m_new = torch.maximum(m, S.max(dim=1).values)
P = torch.exp(S - m_new[:, None])
l_new = torch.exp(m - m_new) * l + P.sum(dim=1)
O = O * (torch.exp(m - m_new) * l / l_new.clamp_min(1e-12))[:, None] \
+ (P @ Vj) / l_new.clamp_min(1e-12)[:, None]
m, l = m_new, l_new
return O
Tests
T, d = 32, 16
Q = torch.randn(T, d); K = torch.randn(T, d); V = torch.randn(T, d)
ref = F.scaled_dot_product_attention(Q.unsqueeze(0), K.unsqueeze(0), V.unsqueeze(0)).squeeze(0)
ours = flash_attn_naive(Q, K, V, tile=8)
assert torch.allclose(ours, ref, atol=1e-4)
print("FlashAttn illustrative OK")
Problem
Implement a Pre-LN Transformer encoder block: LayerNorm → MHA → residual → LayerNorm → MLP → residual.
import torch.nn as nn
class EncoderBlock(nn.Module):
def __init__(self, d, h, mlp_ratio=4):
super().__init__()
self.ln1 = nn.LayerNorm(d); self.attn = MHA(d, h)
self.ln2 = nn.LayerNorm(d)
self.mlp = nn.Sequential(nn.Linear(d, mlp_ratio*d), nn.GELU(), nn.Linear(mlp_ratio*d, d))
def forward(self, x, mask=None):
x = x + self.attn(self.ln1(x), mask)
x = x + self.mlp(self.ln2(x))
return x
Tests
m = EncoderBlock(64, 4); x = torch.randn(2, 7, 64)
assert m(x).shape == x.shape; print("Encoder block OK")
Problem
Implement a Pre-LN causal decoder block (used by GPT, LLaMA). Self-attention only (no cross-attn).
import torch.nn as nn
class DecoderBlock(nn.Module):
def __init__(self, d, h, mlp_ratio=4):
super().__init__()
self.ln1 = nn.LayerNorm(d); self.attn = MHA(d, h)
self.ln2 = nn.LayerNorm(d)
self.mlp = nn.Sequential(nn.Linear(d, mlp_ratio*d), nn.GELU(), nn.Linear(mlp_ratio*d, d))
def forward(self, x):
T = x.size(1); m = causal_mask(T, x.device)
x = x + self.attn(self.ln1(x), m)
x = x + self.mlp(self.ln2(x))
return x
Tests
m = DecoderBlock(32, 4); x = torch.randn(1, 6, 32)
assert m(x).shape == x.shape; print("Decoder block OK")
Problem
Train a tiny BPE on a list of words: count adjacent symbol pairs, repeatedly merge the most frequent into a new token.
from collections import Counter
def get_pairs(word_freqs):
pairs = Counter()
for word, f in word_freqs.items():
sym = word.split()
for i in range(len(sym)-1):
pairs[(sym[i], sym[i+1])] += f
return pairs
def merge(word_freqs, pair):
bigram = ' '.join(pair); rep = ''.join(pair)
return {w.replace(bigram, rep): f for w, f in word_freqs.items()}
def train_bpe(corpus, num_merges):
# corpus: list of strings; treat each word as space-separated chars + </w>
word_freqs = Counter([' '.join(list(w)) + ' </w>' for w in corpus])
merges = []
for _ in range(num_merges):
pairs = get_pairs(word_freqs)
if not pairs: break
best = pairs.most_common(1)[0][0]
word_freqs = merge(word_freqs, best)
merges.append(best)
return merges
Tests
corpus = ['low', 'low', 'low', 'low', 'low',
'lower', 'lower',
'newest', 'newest', 'newest', 'newest', 'newest', 'newest',
'widest', 'widest', 'widest']
merges = train_bpe(corpus, num_merges=10)
assert ('e', 's') in merges or ('es', 't') in merges
print("BPE OK", merges[:5])
Problem
Given logits, return a single sampled token using temperature scaling, then top-k, then top-p (nucleus) filtering.
import torch
def sample_token(logits, temperature=1.0, top_k=0, top_p=0.0):
logits = logits.clone() / max(temperature, 1e-8)
if top_k > 0:
v, _ = torch.topk(logits, top_k)
logits[logits < v[..., -1, None]] = -float('inf')
if top_p > 0.0:
sorted_l, idx = logits.sort(descending=True)
cum = torch.softmax(sorted_l, dim=-1).cumsum(dim=-1)
mask = cum > top_p
mask[..., 1:] = mask[..., :-1].clone(); mask[..., 0] = False
sorted_l[mask] = -float('inf')
logits = torch.full_like(logits, -float('inf')).scatter_(-1, idx, sorted_l)
probs = torch.softmax(logits, dim=-1)
return torch.multinomial(probs, num_samples=1)
Tests
torch.manual_seed(0)
logits = torch.tensor([[1., 2., 3., 4., 5.]])
tok = sample_token(logits, temperature=0.001) # near-argmax
assert tok.item() == 4
tok = sample_token(logits, top_k=1, temperature=1.0)
assert tok.item() == 4
print("sampling OK")
Problem
Implement beam search of width B for an autoregressive scorer. Keep cumulative log-probs.
import torch
def beam_search(score_fn, init_seq, beam=4, max_len=10, eos=2):
# score_fn(seq) -> log-probs over next token (V,)
beams = [(init_seq, 0.0)]
for _ in range(max_len):
cands = []
for seq, lp in beams:
if seq[-1] == eos:
cands.append((seq, lp)); continue
log_p = score_fn(seq)
top_lp, top_id = torch.topk(log_p, beam)
for k in range(beam):
cands.append((seq + [int(top_id[k])], lp + float(top_lp[k])))
beams = sorted(cands, key=lambda x: x[1], reverse=True)[:beam]
if all(s[-1] == eos for s, _ in beams):
break
return beams
Tests
def fake_score(seq): # toy: prefer 1, then 2 (eos)
out = torch.full((5,), -10.0)
out[1] = 0.0; out[2] = -1.0
return out
out = beam_search(fake_score, [0], beam=3, max_len=5, eos=2)
assert out[0][0][:3] == [0, 1, 1]; print("beam OK", out[0])
Advanced implementation. All beams advance in ONE forward per step (beams live in the batch dim); selection is one top-k over the flattened (beam, vocab) grid. The -inf initial scores prevent duplicate start-beams winning top-k B times — the classic batched-beam bug. Verified: same sequence and score as sequential beam search; ~6x faster at width 8.
import torch
import torch.nn.functional as F
@torch.no_grad()
def beam_search_batched(model, idx, max_new, beam_width):
seqs = idx.repeat(beam_width, 1)
scores = torch.full((beam_width,), float("-inf")); scores[0] = 0.0
for _ in range(max_new):
logp = F.log_softmax(model(seqs)[:, -1], dim=-1) # ONE forward for all beams
V = logp.size(-1)
top = (scores[:, None] + logp).view(-1).topk(beam_width)
seqs = torch.cat([seqs[top.indices // V], (top.indices % V)[:, None]], 1)
scores = top.values
best = scores.argmax()
return seqs[best:best + 1], scores[best].item()
Problem
Given A ∈ RN×D and B ∈ RM×D, compute the N ×M cosine-similarity matrix.
def cos_sim(A, B):
A = A / (A.norm(dim=-1, keepdim=True) + 1e-12)
B = B / (B.norm(dim=-1, keepdim=True) + 1e-12)
return A @ B.T
Tests
A = torch.randn(4, 8); B = torch.randn(5, 8)
S = cos_sim(A, B); assert S.shape == (4, 5) and (S.abs() <= 1.0 + 1e-6).all()
print("cos_sim OK")
Problem
Implement brute-force kNN over a query batch and a key set with cosine distance.
import torch
def knn(queries, keys, k=5):
sim = cos_sim(queries, keys) # higher = closer
vals, idx = torch.topk(sim, k, dim=-1)
return idx, vals
Tests
torch.manual_seed(0)
keys = torch.randn(100, 16)
q = keys[0:1] + 0.01 * torch.randn(1, 16) # query close to key 0
idx, _ = knn(q, keys, k=5); assert 0 in idx[0].tolist()
print("knn OK")
Problem
Given two views’ embeddings z1, z2 ∈ RN×D, compute the symmetric InfoNCE loss (temperature τ).
import torch
import torch.nn.functional as F
def info_nce(z1, z2, tau=0.1):
z = torch.cat([z1, z2], dim=0) # (2N, D)
z = z / z.norm(dim=-1, keepdim=True)
sim = (z @ z.T) / tau # (2N, 2N)
N = z1.size(0)
sim.fill_diagonal_(-1e9)
pos = torch.cat([torch.arange(N, 2*N), torch.arange(0, N)]).to(z.device)
return F.cross_entropy(sim, pos)
Tests
z = torch.randn(16, 32); z2 = z + 0.001 * torch.randn_like(z)
loss = info_nce(z, z2, tau=0.07).item()
assert loss < 1.0 # pairs are nearly identical
print("info_nce OK", loss)
Problem
Implement the triplet margin loss max(0, ∥a −p∥−∥a −n∥+ m) for a batch.
import torch
def triplet_loss(a, p, n, margin=0.2):
dap = (a - p).norm(dim=-1)
dan = (a - n).norm(dim=-1)
return torch.clamp_min(dap - dan + margin, 0.0).mean()
Tests
a = torch.zeros(4, 8); p = torch.zeros(4, 8) + 0.1; n = torch.ones(4, 8)
assert triplet_loss(a, p, n, margin=0.2) < 1.0; print("triplet OK")
Problem
Compute IoU between two batches of axis-aligned boxes in [x1, y1, x2, y2] form.
import torch
def iou(a, b):
# a: (N, 4), b: (M, 4). returns (N, M)
inter_x1 = torch.max(a[:, None, 0], b[None, :, 0])
inter_y1 = torch.max(a[:, None, 1], b[None, :, 1])
inter_x2 = torch.min(a[:, None, 2], b[None, :, 2])
inter_y2 = torch.min(a[:, None, 3], b[None, :, 3])
iw = (inter_x2 - inter_x1).clamp(min=0)
ih = (inter_y2 - inter_y1).clamp(min=0)
inter = iw * ih
area_a = (a[:, 2] - a[:, 0]) * (a[:, 3] - a[:, 1])
area_b = (b[:, 2] - b[:, 0]) * (b[:, 3] - b[:, 1])
union = area_a[:, None] + area_b[None, :] - inter
return inter / union.clamp(min=1e-9)
Tests
a = torch.tensor([[0., 0., 10., 10.]])
b = torch.tensor([[5., 5., 15., 15.], [0., 0., 10., 10.]])
out = iou(a, b)
assert abs(out[0,0].item() - 25/175) < 1e-6 and abs(out[0,1].item() - 1.0) < 1e-6
print("iou OK")
Problem
Standard greedy NMS: sort by score, iteratively keep top, suppress others whose IoU exceeds thresh.
def nms(boxes, scores, thresh=0.5):
order = scores.argsort(descending=True).tolist()
keep = []
while order:
i = order.pop(0); keep.append(i)
if not order: break
ious = iou(boxes[i:i+1], boxes[order]).squeeze(0)
order = [j for j, iv in zip(order, ious.tolist()) if iv <= thresh]
return keep
Tests
B = torch.tensor([[0,0,10,10], [1,1,11,11], [50,50,60,60]], dtype=torch.float32)
S = torch.tensor([0.9, 0.8, 0.7])
keep = nms(B, S, thresh=0.5)
assert keep == [0, 2]; print("nms OK")
Problem
Match A anchors to G ground-truth boxes by IoU with thresholds for positive (≥0.5), negative (< 0.4), and ignore (in between).
import torch
def match_anchors(anchors, gt, pos_thr=0.5, neg_thr=0.4):
if len(gt) == 0:
return torch.full((len(anchors),), -1, dtype=torch.long) # all neg
ious = iou(anchors, gt) # (A, G)
max_iou, gt_idx = ious.max(dim=1)
labels = torch.full((len(anchors),), -2, dtype=torch.long) # -2 = ignore
labels[max_iou < neg_thr] = -1 # negative
labels[max_iou >= pos_thr] = gt_idx[max_iou >= pos_thr] # positive (idx into gt)
# ensure each gt has at least one positive
best_anchor = ious.argmax(dim=0)
labels[best_anchor] = torch.arange(len(gt))
return labels
Tests
anchors = torch.tensor([[0,0,10,10], [10,10,20,20], [0,0,5,5]], dtype=torch.float32)
gt = torch.tensor([[0,0,10,10]], dtype=torch.float32)
lab = match_anchors(anchors, gt)
assert lab[0].item() == 0 and lab[1].item() == -1
print("anchor matching OK", lab.tolist())
Problem
Implement PatchEmbed: split image into non-overlapping patches and project each to a D-dim embedding via Conv2d.
import torch.nn as nn
class PatchEmbed(nn.Module):
def __init__(self, img=224, patch=16, dim=768, in_ch=3):
super().__init__()
self.proj = nn.Conv2d(in_ch, dim, kernel_size=patch, stride=patch)
def forward(self, x): # (B, C, H, W)
x = self.proj(x) # (B, D, H/p, W/p)
return x.flatten(2).transpose(1, 2) # (B, N, D)
Tests
m = PatchEmbed(); y = m(torch.randn(2, 3, 224, 224))
assert y.shape == (2, 196, 768); print("patch embed OK")
Problem
Standard ResNet bottleneck: 1×1 → 3×3 → 1×1 with stride and a projection shortcut when channels or stride change.
import torch.nn as nn
import torch.nn.functional as F
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, in_c, planes, stride=1):
super().__init__()
out_c = planes * self.expansion
self.c1 = nn.Conv2d(in_c, planes, 1, bias=False); self.bn1 = nn.BatchNorm2d(planes)
self.c2 = nn.Conv2d(planes, planes, 3, stride, 1, bias=False); self.bn2 = nn.BatchNorm2d(planes)
self.c3 = nn.Conv2d(planes, out_c, 1, bias=False); self.bn3 = nn.BatchNorm2d(out_c)
self.short = nn.Sequential() if (stride==1 and in_c==out_c) else nn.Sequential(
nn.Conv2d(in_c, out_c, 1, stride, bias=False), nn.BatchNorm2d(out_c))
def forward(self, x):
h = F.relu(self.bn1(self.c1(x)))
h = F.relu(self.bn2(self.c2(h)))
h = self.bn3(self.c3(h))
return F.relu(h + self.short(x))
Tests
m = Bottleneck(64, 64, stride=2); y = m(torch.randn(1, 64, 32, 32))
assert y.shape == (1, 256, 16, 16); print("bottleneck OK")
Problem
Implement focal loss for binary classification: −α(1 −pt)γ log pt.
import torch
import torch.nn.functional as F
def focal_loss(logits, targets, alpha=0.25, gamma=2.0):
p = torch.sigmoid(logits)
pt = torch.where(targets == 1, p, 1 - p)
a = torch.where(targets == 1, torch.tensor(alpha), torch.tensor(1 - alpha))
bce = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')
return (a * (1 - pt) ** gamma * bce).mean()
Tests
z = torch.randn(100); y = (torch.rand(100) > 0.95).float()
assert focal_loss(z, y) >= 0; print("focal OK")
Problem
Soft Dice loss for binary segmentation: 1 −2|A∩B|+ε |A|+|B|+ε.
def dice_loss(probs, targets, eps=1e-6):
p = probs.flatten(1); t = targets.flatten(1)
inter = (p * t).sum(dim=1)
return (1 - (2 * inter + eps) / (p.sum(dim=1) + t.sum(dim=1) + eps)).mean()
Tests
p = torch.zeros(2, 1, 4, 4); t = torch.zeros_like(p)
assert dice_loss(torch.ones_like(p), t) > 0.9
assert dice_loss(t, t) < 1e-5; print("dice OK")
Problem
Implement label-smoothed CE with smoothing ε: target distribution becomes (1 −ε)δy + ε/C.
import torch.nn.functional as F
def smooth_ce(logits, targets, eps=0.1):
C = logits.size(-1)
log_p = F.log_softmax(logits, dim=-1)
nll = -log_p.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
smooth = -log_p.mean(dim=-1)
return ((1 - eps) * nll + eps * smooth).mean()
Tests
z = torch.randn(8, 10); y = torch.randint(0, 10, (8,))
assert smooth_ce(z, y, eps=0.0).item() == F.cross_entropy(z, y).item()
print("smooth CE OK")
Problem
Implement KL(P∥Q) for discrete distributions, both with explicit probabilities and from logits.
import torch.nn.functional as F
def kl_from_logits(p_logits, q_logits):
log_p = F.log_softmax(p_logits, dim=-1)
log_q = F.log_softmax(q_logits, dim=-1)
return (log_p.exp() * (log_p - log_q)).sum(dim=-1).mean()
Tests
z = torch.randn(4, 6)
assert kl_from_logits(z, z).item() < 1e-6
print("KL OK")
Problem
Implement the forward-only CTC log-likelihood for a single sequence (no batching, no GPU). Reference: Graves 2006.
import math
def ctc_log_loss(log_probs, target, blank=0):
# log_probs: (T, V), target: list[int] without blanks
T, V = log_probs.shape
L = 2 * len(target) + 1
ext = [blank] + sum(([t, blank] for t in target), [])
INF = -float('inf')
a = [[INF]*L for _ in range(T)]
a[0][0] = log_probs[0, blank]
if L > 1: a[0][1] = log_probs[0, ext[1]]
for t in range(1, T):
for s in range(L):
cands = [a[t-1][s]]
if s-1 >= 0: cands.append(a[t-1][s-1])
if s-2 >= 0 and ext[s] != blank and ext[s] != ext[s-2]:
cands.append(a[t-1][s-2])
a[t][s] = max(cands) + math.log(sum(math.exp(c - max(cands)) for c in cands)) + log_probs[t, ext[s]]
return -max(a[-1][L-1], a[-1][L-2])
Tests
# Sanity: target == argmax sequence should have low loss
T, V = 5, 4
lp = torch.full((T, V), -10.0); lp[:, 1] = 0.0
loss = ctc_log_loss(lp.numpy(), [1, 1, 1])
assert loss < 5.0; print("CTC OK", loss)
Problem
Given a β schedule of length T, sample xt in closed form: xt = √ᾱt x0 + √1 −ᾱt ε with ε ∼N(0, I).
import torch
def make_schedule(T, beta1=1e-4, beta2=0.02):
betas = torch.linspace(beta1, beta2, T)
alphas = 1.0 - betas
abar = torch.cumprod(alphas, dim=0)
return betas, alphas, abar
def q_sample(x0, t, abar, eps=None):
eps = torch.randn_like(x0) if eps is None else eps
a = abar[t].view(-1, *([1] * (x0.dim() - 1)))
return a.sqrt() * x0 + (1 - a).sqrt() * eps, eps
Tests
T = 1000; _, _, abar = make_schedule(T)
x0 = torch.randn(2, 3, 32, 32); t = torch.tensor([0, T-1])
xt, _ = q_sample(x0, t, abar)
assert torch.allclose(xt[0], x0[0], atol=1e-2)
assert xt[1].std().item() > 0.5 # second sample is essentially noise
print("DDPM forward OK")
Problem
Implement the DDPM training loss Ex0,t,ε ∥ε −εθ(xt, t)∥2.
import torch
import torch.nn.functional as F
def ddpm_loss(model, x0, abar):
t = torch.randint(0, abar.size(0), (x0.size(0),), device=x0.device)
xt, eps = q_sample(x0, t, abar)
pred = model(xt, t)
return F.mse_loss(pred, eps)
Tests
class Tiny(nn.Module):
def __init__(self): super().__init__(); self.c = nn.Conv2d(3, 3, 3, padding=1)
def forward(self, x, t): return self.c(x)
m = Tiny(); _, _, abar = make_schedule(50)
loss = ddpm_loss(m, torch.randn(2, 3, 8, 8), abar); assert loss > 0
print("DDPM loss OK")
Problem
One reverse step: from xt and εθ(xt, t) compute xt−1.
import torch
def ddpm_step(xt, t, eps_pred, betas, alphas, abar):
a = alphas[t]; b = betas[t]; ab = abar[t]
coef = (1 - a) / (1 - ab).sqrt()
mean = (1.0 / a.sqrt()) * (xt - coef * eps_pred)
if t == 0:
return mean
noise = torch.randn_like(xt)
return mean + b.sqrt() * noise
Tests
T = 100; betas, alphas, abar = make_schedule(T)
xt = torch.randn(1, 1, 4, 4); eps = torch.zeros_like(xt)
out = ddpm_step(xt, t=10, eps_pred=eps, betas=betas, alphas=alphas, abar=abar)
assert out.shape == xt.shape; print("DDPM step OK")
Problem
Combine conditional and unconditional model outputs: ε˜ = (1 + w)εc −wεu.
def cfg(eps_uncond, eps_cond, w=7.5):
return (1 + w) * eps_cond - w * eps_uncond
Tests
eu = torch.zeros(4); ec = torch.ones(4)
assert torch.allclose(cfg(eu, ec, w=0.0), ec)
assert torch.allclose(cfg(eu, ec, w=1.0), 2 * ec)
print("CFG OK")
Problem
Conditional flow matching: minimize Et,x0,x1 ∥vθ(xt, t) −(x1 −x0)∥2 where xt = (1 −t)x0 + tx1, t ∼U(0, 1).
import torch
import torch.nn.functional as F
def flow_matching_loss(model, x0, x1):
t = torch.rand(x0.size(0), device=x0.device).view(-1, *([1]*(x0.dim()-1)))
xt = (1 - t) * x0 + t * x1
target = x1 - x0
return F.mse_loss(model(xt, t.flatten()), target)
Tests
class TinyV(nn.Module):
def __init__(self): super().__init__(); self.lin = nn.Linear(4, 4)
def forward(self, x, t): return self.lin(x.flatten(1)).view_as(x)
m = TinyV()
loss = flow_matching_loss(m, torch.randn(8, 1, 2, 2), torch.randn(8, 1, 2, 2))
assert loss > 0; print("FM OK")
Problem
Compute the REINFORCE policy-gradient loss given log-probabilities and (un-/baselined) returns.
def reinforce_loss(log_probs, returns, baseline=None):
if baseline is not None: returns = returns - baseline
return -(log_probs * returns).mean()
Tests
lp = torch.log(torch.tensor([0.1, 0.5, 0.4]))
R = torch.tensor([1., 0., -1.])
assert reinforce_loss(lp, R).item() != 0; print("reinforce OK")
Problem
Given new and old log-probs and advantages, implement the PPO clipped objective with ε = 0.2.
import torch
def ppo_loss(logp_new, logp_old, adv, clip=0.2):
ratio = (logp_new - logp_old).exp()
surr1 = ratio * adv
surr2 = ratio.clamp(1 - clip, 1 + clip) * adv
return -torch.min(surr1, surr2).mean()
Tests
lpn = torch.tensor([0.0, 0.0]); lpo = torch.tensor([0.0, 0.0])
adv = torch.tensor([1.0, -1.0])
assert ppo_loss(lpn, lpo, adv).item() == 0.0 # ratio = 1, both surr identical
print("PPO OK")
Problem
Implement GAE(γ, λ) given rewards and value estimates.
def gae(rewards, values, gamma=0.99, lam=0.95, last_value=0.0):
T = len(rewards)
adv = [0.0] * T
g = 0.0
next_v = last_value
for t in reversed(range(T)):
delta = rewards[t] + gamma * next_v - values[t]
g = delta + gamma * lam * g
adv[t] = g
next_v = values[t]
return adv
Tests
adv = gae([1, 1, 1], [0, 0, 0], gamma=1.0, lam=1.0)
assert adv == [3.0, 2.0, 1.0]; print("GAE OK")
Problem
Group Relative Policy Optimization computes advantages by subtracting the group mean and dividing by the group std. Used in DeepSeek-R1 / R1-Zero training.
def grpo_advantage(rewards, eps=1e-8):
# rewards: (G,) per group
mu = rewards.mean(); sd = rewards.std()
return (rewards - mu) / (sd + eps)
Tests
A = grpo_advantage(torch.tensor([1.0, 2.0, 3.0, 4.0]))
assert abs(A.mean().item()) < 1e-6; print("GRPO OK", A.tolist())
Problem
r , where A ∈ Rr×din, B ∈ Rdout×r.Implement a LoRA adapter over a frozen nn.Linear: h = Wx + (BA)x · α
import math
import torch
import torch.nn as nn
class LoRALinear(nn.Module):
def __init__(self, base: nn.Linear, r=8, alpha=16):
super().__init__()
self.base = base
for p in self.base.parameters(): p.requires_grad = False
self.A = nn.Parameter(torch.zeros(r, base.in_features))
self.B = nn.Parameter(torch.zeros(base.out_features, r))
nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
self.scale = alpha / r
def forward(self, x):
return self.base(x) + (x @ self.A.T) @ self.B.T * self.scale
Tests
base = nn.Linear(16, 16); m = LoRALinear(base, r=4, alpha=8)
y = m(torch.randn(2, 16))
trainable = sum(p.numel() for p in m.parameters() if p.requires_grad)
assert trainable == 4*16 + 16*4; print("LoRA OK")
Problem
Quantize a float tensor to INT8 with symmetric per-tensor scale; dequantize and check error.
import torch
def quantize(x, bits=8):
qmax = 2 ** (bits - 1) - 1
s = x.abs().max() / qmax
q = torch.round(x / s).clamp(-qmax, qmax).to(torch.int8)
return q, s
def dequantize(q, s):
return q.float() * s
Tests
x = torch.randn(1000)
q, s = quantize(x); xh = dequantize(q, s)
err = (x - xh).abs().mean().item()
assert err < 0.02; print("int8 quant OK", err)
Problem
Implement a top-2 sparse MoE router with auxiliary load-balancing loss.
import torch
import torch.nn as nn
import torch.nn.functional as F
class MoE(nn.Module):
def __init__(self, d, n_experts, k=2):
super().__init__()
self.gate = nn.Linear(d, n_experts, bias=False)
self.experts = nn.ModuleList([nn.Linear(d, d) for _ in range(n_experts)])
self.k = k
def forward(self, x):
# x: (B*T, D)
scores = self.gate(x)
topk_v, topk_i = scores.topk(self.k, dim=-1)
weights = F.softmax(topk_v, dim=-1)
out = torch.zeros_like(x)
for k in range(self.k):
for e in range(len(self.experts)):
mask = topk_i[:, k] == e
if mask.any():
out[mask] += weights[mask, k:k+1] * self.experts[e](x[mask])
# aux load balance loss
prob = F.softmax(scores, dim=-1).mean(dim=0)
load = (topk_i == torch.arange(len(self.experts), device=x.device).view(-1,1,1)).any(0).float().mean(0)
aux = (prob * load).sum() * len(self.experts)
return out, aux
Tests
m = MoE(32, n_experts=4, k=2); y, aux = m(torch.randn(8, 32))
assert y.shape == (8, 32) and aux.item() >= 0
print("MoE OK")
Problem
Implement the dynamic-scale loop: scale the loss, unscale grads, skip on inf/nan, otherwise update and possibly grow the scale.
import torch
class Scaler:
def __init__(self, init=2.0**16, growth=2.0, backoff=0.5, growth_int=2000):
self.scale, self.growth, self.backoff, self.gi = init, growth, backoff, growth_int
self.steps_no_skip = 0
def scale_loss(self, loss): return loss * self.scale
def step(self, params_grads, optimizer_step):
bad = any(not torch.isfinite(g).all() for _, g in params_grads)
if bad:
self.scale *= self.backoff; self.steps_no_skip = 0; return False
for p, g in params_grads:
g.div_(self.scale)
optimizer_step()
self.steps_no_skip += 1
if self.steps_no_skip % self.gi == 0:
self.scale *= self.growth
return True
Tests
s = Scaler(init=8.0)
p = torch.zeros(3); g = torch.tensor([float('inf'),0.0,0.0])
ok = s.step([(p, g)], lambda: None)
assert ok is False and s.scale == 4.0
print("Scaler OK")
Problem
Implement a tensor Dataset and a basic DataLoader (batching, shuffling).
from torch.utils.data import Dataset, DataLoader
class TensorDS(Dataset):
def __init__(self, X, y): self.X, self.y = X, y
def __len__(self): return len(self.X)
def __getitem__(self, i): return self.X[i], self.y[i]
def make_loader(X, y, batch=32, shuffle=True):
return DataLoader(TensorDS(X, y), batch_size=batch, shuffle=shuffle)
Tests
X = torch.arange(100).float().view(-1, 1); y = torch.arange(100)
loader = make_loader(X, y, batch=8, shuffle=False)
xb, yb = next(iter(loader))
assert xb.shape == (8, 1) and yb[0].item() == 0; print("DataLoader OK")
Problem
Given a list of variable-length 1-D tensors, pad to the max length and return a key-padding mask.
import torch
def pad_collate(batch, pad_id=0):
lens = [len(t) for t in batch]
L = max(lens)
out = torch.full((len(batch), L), pad_id, dtype=batch[0].dtype)
for i, t in enumerate(batch): out[i, :len(t)] = t
mask = torch.arange(L)[None, :] >= torch.tensor(lens)[:, None] # True = pad
return out, mask
Tests
seqs = [torch.tensor([1,2,3]), torch.tensor([4,5]), torch.tensor([6])]
out, mask = pad_collate(seqs, pad_id=-1)
assert out.tolist() == [[1,2,3],[4,5,-1],[6,-1,-1]]
assert mask[2].tolist() == [False, True, True]; print("pad OK")
Problem
Implement the original Transformer’s sinusoidal PE: even dims use sin, odd dims use cos.
import math
import torch
def sinusoidal_pe(L, D):
pe = torch.zeros(L, D)
pos = torch.arange(L).unsqueeze(1).float()
div = torch.exp(torch.arange(0, D, 2).float() * (-math.log(10000.0) / D))
pe[:, 0::2] = torch.sin(pos * div)
pe[:, 1::2] = torch.cos(pos * div)
return pe
Tests
pe = sinusoidal_pe(10, 16)
assert pe.shape == (10, 16) and torch.allclose(pe[0, 1::2], torch.ones(8))
print("PE OK")
Problem
Compare a function’s analytical gradient (autograd) against the numerical gradient via central differences.
import torch
def grad_check(fn, x, eps=1e-5):
# fn: tensor -> scalar
x = x.clone().detach().requires_grad_(True)
fn(x).backward()
g_an = x.grad.clone()
g_num = torch.zeros_like(x)
for i in range(x.numel()):
e = torch.zeros_like(x); e.view(-1)[i] = eps
g_num.view(-1)[i] = (fn(x + e).item() - fn(x - e).item()) / (2 * eps)
return (g_an - g_num).abs().max().item()
Tests
err = grad_check(lambda z: (z ** 2).sum(), torch.randn(5))
assert err < 1e-4; print("grad check OK", err)
Problem
Implement vanilla K-Means with k-means++ initialisation.
import numpy as np
def kmeanspp_init(X, k, rng):
idx = [rng.randint(len(X))]
for _ in range(1, k):
d2 = np.min(((X[:, None] - X[idx][None]) ** 2).sum(-1), axis=1)
p = d2 / d2.sum()
idx.append(rng.choice(len(X), p=p))
return X[idx]
def kmeans(X, k=4, iters=20, seed=0):
rng = np.random.default_rng(seed)
C = kmeanspp_init(X, k, rng)
for _ in range(iters):
d = ((X[:, None] - C[None]) ** 2).sum(-1)
a = d.argmin(axis=1)
C = np.stack([X[a == j].mean(0) if (a==j).any() else C[j] for j in range(k)])
return C, a
Tests
X = np.concatenate([np.random.randn(50, 2) + np.array([-5, 0]),
np.random.randn(50, 2) + np.array([5, 0])])
C, a = kmeans(X, k=2)
assert len(set(a)) == 2; print("kmeans OK")
Advanced implementation. Avoid the (n, k, d) broadcast tensor with the quadratic expansion (distances in (n, k)), and replace the per-cluster centroid loop with one one-hot GEMM. Verified: identical trajectory and inertia to the loop version; ~16x less peak memory at n=20k.
import numpy as np
def kmeans_step(X, C):
d2 = (X**2).sum(1)[:, None] - 2 * X @ C.T + (C**2).sum(1)[None, :] # (n, k)
assign = d2.argmin(1)
onehot = (assign[:, None] == np.arange(len(C))).astype(X.dtype)
counts = onehot.sum(0)
newC = np.where(counts[:, None] > 0,
(onehot.T @ X) / np.maximum(counts, 1)[:, None], C) # one GEMM
return newC, assign
Problem
Train binary logistic regression with BCE and SGD. Report accuracy.
import numpy as np
def logreg_fit(X, y, lr=0.1, epochs=200):
w = np.zeros(X.shape[1]); b = 0.0
for _ in range(epochs):
z = X @ w + b
p = 1.0 / (1.0 + np.exp(-z))
gw = X.T @ (p - y) / len(y)
gb = (p - y).mean()
w -= lr * gw; b -= lr * gb
return w, b
Tests
np.random.seed(0)
X = np.random.randn(200, 3)
y = (X @ np.array([1.0, -1.0, 0.5]) > 0).astype(np.float64)
w, b = logreg_fit(X, y); pred = (X @ w + b > 0).astype(np.float64)
assert (pred == y).mean() > 0.9; print("logreg OK")
Problem
Implement PCA on a centered data matrix using thin SVD; project to k components.
import numpy as np
def pca(X, k):
Xc = X - X.mean(axis=0)
U, S, Vt = np.linalg.svd(Xc, full_matrices=False)
return Xc @ Vt[:k].T, Vt[:k]
Tests
X = np.random.randn(100, 6) @ np.random.randn(6, 6)
Z, V = pca(X, k=2)
assert Z.shape == (100, 2); print("pca OK")
Problem
Implement accuracy(logits, targets, k) that returns top-k classification accuracy.
def accuracy(logits, targets, k=1):
topk = logits.topk(k, dim=-1).indices # (N, k)
return (topk == targets.unsqueeze(-1)).any(dim=-1).float().mean().item()
Tests
z = torch.tensor([[0.1, 0.9], [0.8, 0.2]]); y = torch.tensor([1, 0])
assert accuracy(z, y, k=1) == 1.0; print("acc OK")
Problem
Sample from a Categorical(π) in a differentiable way.
import torch
import torch.nn.functional as F
def gumbel_softmax(logits, tau=1.0, hard=False):
g = -torch.log(-torch.log(torch.rand_like(logits) + 1e-20) + 1e-20)
y = F.softmax((logits + g) / tau, dim=-1)
if hard:
idx = y.argmax(dim=-1, keepdim=True)
y_hard = torch.zeros_like(y).scatter_(-1, idx, 1.0)
y = (y_hard - y).detach() + y # straight-through
return y
Tests
torch.manual_seed(0)
y = gumbel_softmax(torch.tensor([[1.0, 2.0, 3.0]]).expand(1000, -1), tau=0.5)
assert abs(y.argmax(-1).float().mean().item() - 1.6) < 0.5 # mostly index 2
print("gumbel OK")
Problem
Compute the VAE objective: BCE reconstruction + KL of q(z|x) ∼N(µ, σ2) to N(0, I).
import torch.nn.functional as F
def vae_loss(x_hat, x, mu, logvar):
rec = F.binary_cross_entropy(x_hat, x, reduction='sum')
kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum()
return (rec + kl) / x.size(0)
Tests
x = torch.rand(4, 8); xh = torch.rand(4, 8)
mu = torch.zeros(4, 2); lv = torch.zeros(4, 2)
loss = vae_loss(xh, x, mu, lv); assert loss > 0; print("VAE ELBO OK")
Problem
Compute the per-token cross-entropy of a Transformer decoder under teacher forcing, ignoring padded posi- tions.
import torch.nn.functional as F
def lm_loss(logits, targets, pad_id=-100):
# logits: (B, T, V), targets: (B, T)
B, T, V = logits.shape
return F.cross_entropy(logits.reshape(-1, V), targets.reshape(-1), ignore_index=pad_id)
Tests
z = torch.randn(2, 4, 5)
t = torch.tensor([[1, 2, 3, -100], [0, 1, -100, -100]])
assert lm_loss(z, t).item() > 0; print("LM loss OK")
Problem
Build a token-embedding layer whose weights are tied to the output classifier (used by GPT etc.).
import torch.nn as nn
class TiedHeadLM(nn.Module):
def __init__(self, vocab, d):
super().__init__()
self.emb = nn.Embedding(vocab, d)
# no separate output Linear; reuse self.emb.weight
def forward(self, ids):
h = self.emb(ids)
return h, h @ self.emb.weight.T # tied logits
Tests
m = TiedHeadLM(100, 32)
ids = torch.randint(0, 100, (2, 5))
h, logits = m(ids); assert logits.shape == (2, 5, 100); print("tied head OK")
Problem
Wrap a sequential block so that its forward activations are not stored, but recomputed during backward.
import torch.utils.checkpoint as cp
class CheckpointBlock(nn.Module):
def __init__(self, block): super().__init__(); self.block = block
def forward(self, x): return cp.checkpoint(self.block, x, use_reentrant=False)
Tests
b = nn.Sequential(nn.Linear(8, 8), nn.ReLU(), nn.Linear(8, 8))
m = CheckpointBlock(b); x = torch.randn(2, 8, requires_grad=True)
m(x).sum().backward(); print("checkpoint OK")
Problem
Implement the multi-class margin (Crammer–Singer) loss.
def hinge_multiclass(logits, targets, margin=1.0):
correct = logits.gather(1, targets.unsqueeze(1))
margins = (logits - correct + margin).clamp(min=0)
margins.scatter_(1, targets.unsqueeze(1), 0.0)
return margins.sum(dim=1).mean()
Tests
z = torch.tensor([[2., 0., 0.], [0., 1., 0.]])
y = torch.tensor([0, 1])
out = hinge_multiclass(z, y); assert out.item() == 0.0
print("hinge OK")
Problem
Combine soft-label KL between teacher and student with the regular cross-entropy on hard labels.
import torch.nn.functional as F
def kd_loss(student_logits, teacher_logits, targets, T=4.0, alpha=0.5):
soft = F.kl_div(F.log_softmax(student_logits / T, dim=-1),
F.softmax(teacher_logits / T, dim=-1),
reduction='batchmean') * (T ** 2)
hard = F.cross_entropy(student_logits, targets)
return alpha * soft + (1 - alpha) * hard
Tests
s = torch.randn(4, 10); t = s.detach() + 0.01 * torch.randn_like(s)
y = torch.randint(0, 10, (4,))
assert kd_loss(s, t, y).item() > 0; print("KD OK")
Problem
Implement 1 - cos(a, b) as a per-pair loss with optional reduction.
import torch.nn.functional as F
def cosine_loss(a, b):
return (1 - F.cosine_similarity(a, b, dim=-1)).mean()
Tests
a = torch.randn(4, 8); b = a.clone()
assert cosine_loss(a, b).item() < 1e-6
b = -a; assert abs(cosine_loss(a, b).item() - 2.0) < 1e-5
print("cos loss OK")
Problem
Implement the smooth L1 / Huber loss with delta δ.
import torch
def smooth_l1(x, y, beta=1.0):
d = (x - y).abs()
return torch.where(d < beta, 0.5 * d ** 2 / beta, d - 0.5 * beta).mean()
Tests
out = smooth_l1(torch.tensor([0.0]), torch.tensor([0.5]))
assert abs(out.item() - 0.125) < 1e-6; print("smoothL1 OK")
Problem
Implement a function that clips a list of grad tensors to a global L2 norm c.
import math
def clip_grad_norm(grads, max_norm):
total = math.sqrt(sum(g.pow(2).sum().item() for g in grads if g is not None))
coef = max_norm / (total + 1e-6)
if coef < 1:
for g in grads:
if g is not None: g.mul_(coef)
return total
Tests
g1 = torch.ones(10) * 100.0; g2 = torch.ones(10) * 100.0
n = clip_grad_norm([g1, g2], max_norm=1.0)
assert math.sqrt((g1.pow(2).sum() + g2.pow(2).sum()).item()) < 1.01
print("clip OK")
Problem
Streaming mean/variance over batches without storing all samples (Welford’s algorithm).
class Welford:
def __init__(self): self.n=0; self.mean=0.0; self.M2=0.0
def update(self, x):
self.n += 1; d = x - self.mean
self.mean += d / self.n
self.M2 += d * (x - self.mean)
@property
def var(self): return self.M2 / max(1, self.n - 1)
Tests
data = np.random.randn(1000)
w = Welford()
for v in data: w.update(v)
assert abs(w.mean - data.mean()) < 1e-10
assert abs(w.var - data.var(ddof=1)) < 1e-8
print("welford OK")
Problem
Update running mean/var with momentum m each batch.
def update_running(stats, batch_mean, batch_var, m=0.1):
stats['mean'] = (1 - m) * stats['mean'] + m * batch_mean
stats['var'] = (1 - m) * stats['var'] + m * batch_var
return stats
Tests
s = {'mean': np.zeros(3), 'var': np.ones(3)}
for _ in range(50):
update_running(s, np.array([5, -5, 0.0]), np.array([1, 2, 0.5]))
assert np.allclose(s['mean'], [5, -5, 0], atol=1e-2); print("running OK")
Problem
Compute the diagonal Fisher information for a model (used in EWC continual learning).
import torch
import torch.nn.functional as F
def diag_fisher(model, loader, device='cpu'):
fisher = {n: torch.zeros_like(p) for n, p in model.named_parameters() if p.requires_grad}
model.eval()
for x, y in loader:
model.zero_grad()
loss = F.cross_entropy(model(x.to(device)), y.to(device))
loss.backward()
for n, p in model.named_parameters():
if p.grad is not None:
fisher[n] += p.grad.pow(2)
for n in fisher: fisher[n] /= len(loader)
return fisher
Tests
m = nn.Linear(4, 3)
data = [(torch.randn(2, 4), torch.tensor([0, 1])) for _ in range(3)]
F_diag = diag_fisher(m, data); assert all((v >= 0).all() for v in F_diag.values())
print("Fisher OK")
Problem
Build a sampler that draws from class-balanced indices given imbalanced class labels.
import numpy as np
def class_balanced_indices(labels, n):
labels = np.asarray(labels)
classes, counts = np.unique(labels, return_counts=True)
weights = 1.0 / counts
p = weights[np.searchsorted(classes, labels)]
p /= p.sum()
return np.random.choice(len(labels), size=n, replace=True, p=p)
Tests
y = np.array([0]*100 + [1]*10)
idx = class_balanced_indices(y, 1000)
frac = (y[idx] == 1).mean()
assert 0.4 < frac < 0.6; print("balanced OK", frac)
Problem
Generate a T × T mask that allows each token to attend only within a window of w on each side.
import torch
def sliding_mask(T, w):
i = torch.arange(T)[:, None]
j = torch.arange(T)[None, :]
return torch.where((i - j).abs() <= w, 0.0, float('-inf'))
Tests
M = sliding_mask(5, w=1)
assert M[0, 2] == float('-inf') and M[1, 0] == 0.0; print("sliding OK")
Problem
Implement an in-place, deterministic shuffle (Fisher–Yates) over an index array.
def fisher_yates(idx, rng):
a = list(idx)
for i in range(len(a) - 1, 0, -1):
j = rng.integers(0, i + 1)
a[i], a[j] = a[j], a[i]
return a
Tests
rng = np.random.default_rng(0)
out = fisher_yates(range(10), rng)
assert sorted(out) == list(range(10)) and out != list(range(10))
print("FY shuffle OK")
Problem
Implement input/label mixup: x˜ = λxi + (1 −λ)xj with λ ∼Beta(α, α).
import numpy as np
import torch
import torch.nn.functional as F
def mixup(x, y, alpha=0.2, num_classes=None):
lam = float(np.random.beta(alpha, alpha))
perm = torch.randperm(x.size(0))
x_mix = lam * x + (1 - lam) * x[perm]
if num_classes is not None:
y_oh = F.one_hot(y, num_classes).float()
y_mix = lam * y_oh + (1 - lam) * y_oh[perm]
return x_mix, y_mix
return x_mix, y, y[perm], lam
Tests
x = torch.randn(4, 3, 8, 8); y = torch.randint(0, 5, (4,))
x_m, *_ = mixup(x, y, alpha=0.4)
assert x_m.shape == x.shape; print("mixup OK")
Problem
Implement CutMix: replace a random rectangular region of one image with a region from another.
import numpy as np
import math
import torch
def cutmix(x, y, alpha=1.0):
lam = float(np.random.beta(alpha, alpha))
B, C, H, W = x.shape
perm = torch.randperm(B)
cut_w = int(W * math.sqrt(1 - lam))
cut_h = int(H * math.sqrt(1 - lam))
cx, cy = np.random.randint(W), np.random.randint(H)
x1, y1 = max(cx-cut_w//2, 0), max(cy-cut_h//2, 0)
x2, y2 = min(cx+cut_w//2, W), min(cy+cut_h//2, H)
x_new = x.clone()
x_new[:, :, y1:y2, x1:x2] = x[perm, :, y1:y2, x1:x2]
lam = 1 - ((x2-x1) * (y2-y1) / (W*H))
return x_new, y, y[perm], lam
Tests
x = torch.randn(2, 3, 16, 16); y = torch.tensor([0, 1])
xm, *_ = cutmix(x, y); assert xm.shape == x.shape
print("cutmix OK")
Problem
Given variable-length sequences, pack them into fixed-length blocks separated by EOS tokens.
import torch
def pack_sequences(seqs, block_size, eos=2):
flat = []
for s in seqs:
flat.extend(list(s)); flat.append(eos)
blocks = [flat[i:i+block_size] for i in range(0, len(flat) - block_size + 1, block_size)]
return torch.tensor(blocks)
Tests
out = pack_sequences([[1,2,3], [4,5], [6]], block_size=4, eos=0)
assert out.shape[1] == 4 and out.numel() >= 4; print("pack OK")
Problem
Draft model proposes k tokens; target model verifies. Accept the longest prefix that matches a sample from the target distribution; resample after rejection.
import torch
def speculative_step(draft_probs, target_probs, draft_tokens, rng):
accepted = []
for p_d, p_t, t in zip(draft_probs, target_probs, draft_tokens):
ratio = min(1.0, p_t[t] / max(p_d[t], 1e-12))
if rng.random() < ratio:
accepted.append(int(t))
else:
adj = (p_t - p_d).clamp(min=0); adj /= adj.sum().clamp_min(1e-12)
new_t = int(torch.multinomial(adj, 1))
return accepted + [new_t]
# all accepted: sample one extra from target
extra = int(torch.multinomial(target_probs[-1], 1))
return accepted + [extra]
Tests
import random
rng = random.Random(0)
V = 5; k = 3
dp = [torch.softmax(torch.randn(V), dim=-1) for _ in range(k)]
tp = [torch.softmax(torch.randn(V), dim=-1) for _ in range(k)]
draft = [int(p.argmax()) for p in dp]
acc = speculative_step(dp, tp, draft, rng)
assert 1 <= len(acc) <= k + 1; print("speculative OK", acc)
Problem
Match n predictions to n ground-truth via the (Jonker–Volgenant style) Hungarian algorithm through SciPy.
from scipy.optimize import linear_sum_assignment
def hungarian_match(pred, gt):
cost = ((pred[:, None] - gt[None]) ** 2).sum(-1)
r, c = linear_sum_assignment(cost.numpy() if hasattr(cost, 'numpy') else cost)
return list(zip(r.tolist(), c.tolist()))
Tests
P = torch.tensor([[0,0],[1,0],[0,1]], dtype=torch.float32)
G = torch.tensor([[0,1],[0,0],[1,0]], dtype=torch.float32)
m = hungarian_match(P, G)
assert sorted(m) == [(0,1),(1,2),(2,0)]; print("hungarian OK")
Problem
Normalise an image tensor by ImageNet mean/std and provide its inverse.
import torch
IMNET_MEAN = torch.tensor([0.485, 0.456, 0.406])
IMNET_STD = torch.tensor([0.229, 0.224, 0.225])
def normalize(x): return (x - IMNET_MEAN.view(1,3,1,1)) / IMNET_STD.view(1,3,1,1)
def denormalize(x): return x * IMNET_STD.view(1,3,1,1) + IMNET_MEAN.view(1,3,1,1)
Tests
x = torch.rand(2, 3, 4, 4)
assert torch.allclose(denormalize(normalize(x)), x, atol=1e-6); print("normalize OK")
Notes
Survival tactics for live ML coding:
Always state the shape of every tensor as you go — this catches 80% of bugs without running.
Start with a numerically stable formula (subtract max, log1p, softplus, etc.); robustness to bad inputs scores well.
Walk through edge cases out loud: T = 0, T = 1, batch = 1, zero-length sequence, all-padding, all-masked, single-class, very long sequence.
Prefer explicit shapes over magic. x.view(B, T, H, D) is better than x.view(B, -1, D) if both work.
For PyTorch problems, default to F.scaled_dot_product_attention, F.cross_entropy, F.layer_norm when allowed — they are fast, fused, and correct. Re-implement only if the interviewer asks.
Always end with one or two test cases: random-input shape check, vs.torch reference, and a property-based assertion (e.g. “softmax sums to 1”).
For attention: write the mask first (additive, −∞on disallowed positions), then the QK formula, then the projection. Discuss scaling, causal vs. bidirectional, and KV-cache shape.
For losses: derive the gradient of the loss w.r.t. logits in your head; this exposes whether you actually understand it.
For optimizers: state the per-parameter state size (Adam doubles it). For AdamW, point out the decoupled decay is the only difference.
For diffusion: be prepared to derive the closed-form q(xt|x0) from the recursion in 60 s.
For RL: distinguish on-policy vs. off-policy, advantage estimation, importance ratio, and clip vs. KL penalty.