Exercises — Contributing to open-source ML

Missing step: Claim It (announcing on the issue that you are working on it). The lifecycle in the parent is Find Issue → Claim It → Fork → Branch → Code+Tests → Submit PR → Review → Revisions → Merge. "Claim It" is the step whose entire job is preventing duplicate work — you post a comment so nobody else starts the same fix. Because it was skipped, two people wrote the same patch and the bot flagged the second as a duplicate. This is the only step whose stated purpose matches the failure symptom "duplicate", so it is uniquely determined.

The formula is $Q_{\text{PR}}=C_{\text{correct}}\times C_{\text{clear}}\times C_{\text{tested}}\times C_{\text{justified}}.$ Each factor is binary (0 or 1). Correct ✅, tested ✅, clear ✅ — but the reviewer cannot see the reasoning for the design choice, so $C_{\text{justified}}=0$. Because the factors multiply, one zero collapses everything: $Q_{\text{PR}}=1\times 1\times 1\times 0=0.$ The PR is unmergeable until you add a justification, even though the code works.

RAMP = Read, Active, Makes-work-easier, People-tagged-issues. The rule from the parent's mnemonic: 3 or more "yes" → good fit. Here we have R, A, P as "yes" (3 yeses) and only M as "no". $3\ge 3$, so yes, it is a good fit. The one miss (it doesn't help your daily work) lowers your motivation but does not disqualify it.

Dot product — multiply matching entries and add: $1\cdot 4+2\cdot 5+3\cdot 6=4+10+18=32$. Length of $a$ — square, add, square-root: $\sqrt{1^2+2^2+3^2}=\sqrt{14}$. Length of $b$: $\sqrt{4^2+5^2+6^2}=\sqrt{77}$. $\cos (a,b)=\frac{32}{\sqrt{14}\sqrt{77}}=\frac{32}{\sqrt{1078}}\approx \frac{32}{32.833}\approx \mathrm{0.9746.}$ Close to 1 → the vectors nearly point the same way, which is why the buggy function's answer was fine; only its shape handling was broken.

A 1-D tensor has exactly one axis, numbered $0$. There is no axis $1$ — so .sum(dim=1) asks for a dimension that does not exist, raising RuntimeError: expected 2D. The original code assumed a batch layout [batch, features] where axis $0$ = which example, axis $1$ = the feature entries you sum over. With a bare 1-D vector the features live on axis $0$. So you must sum over dim=0. (The parent's fix instead adds a fake batch axis with unsqueeze(0) turning [3] into [1,3], after which the original dim=1 works — both routes are valid.)

Divide checkpointed by normal — the $B$ and $H$ cancel because they are identical in both: $\frac{{\text{Mem}}_{\text{ckpt}}}{{\text{Mem}}_{\text{normal}}}=\frac{L/k}{L}=\frac{1}{k}=\frac{1}{4}=0.25=25\%.$ So the asymptotic model predicts you keep 25% of activation memory (a 75% saving) at $k=4$. The parent's measured number was 35% kept (65% saved) — real code keeps some non-activation memory (weights, optimizer state) that this simple $O$-model ignores, which is exactly why the parent insisted on empirical benchmarks on top of the derivation.

Memory saved — how much smaller relative to original: $\frac{23.4-4.2}{23.4}=\frac{19.2}{23.4}\approx 0.8205=82.1\%.$ Time added — how much slower relative to original: $\frac{410-180}{180}=\frac{230}{180}\approx 1.2778=127.8\%.$ These match the parent's "82% memory, +128% time" to rounding. Worth it when: the model does not fit at all without checkpointing — a 2.3× slowdown is infinitely better than an out-of-memory crash that gives you no result. If it already fits, the slowdown is pure loss.

• First $L$ — the ordinary forward pass: to produce the output you compute all $L$ layers once. This happens with or without checkpointing.
• Second $L$ — the recomputation during backward. Normally you stored every layer's activation, so backward just reads it (no extra forward work). Checkpointing threw those activations away to save memory, so during backward it must redo the forward pass of the discarded layers to get them back.
• Why $2L$ and not more: you recompute each discarded layer at most once (during its own backward step), so total forward work is the original $L$ plus one extra $L$ of recomputation $=2L$. It is a constant-factor (2×) cost, not a blow-up — that constant factor is precisely the "spare compute cycles" the parent said GPUs can afford.

Root cause: floating-point arithmetic. $\frac{32}{\sqrt{14}\sqrt{77}}$ computed two different ways (once via cosine_similarity, once via manual torch.sqrt) accumulates tiny rounding differences in the last bits, so exact == returns False even though the values agree to ~7 digits. Fix: use torch.allclose(result, expected) instead of == — it checks equality within a tolerance. This is exactly why the parent's test used assert torch.allclose(...). Testing a float for exact equality is the bug, not the math.

## Problem
cosine_similarity crashes on 1-D inputs (RuntimeError: expected 2D)
because it hard-codes .sum(dim=1); 1-D tensors only have axis 0.
Fixes #1234.
 
## Fix
unsqueeze 1-D inputs to shape [1, N] so the existing dim=1 path works,
leaving all higher-rank behaviour byte-identical.
 
## Justification
unsqueeze (vs. hard-coding dim=0) keeps the batched API contract intact,
so 2-D and 3-D callers are unaffected. Covered all input ranks.
 
## Tests
Added test_cosine_similarity_1d comparing against the hand-computed
value 32 / (sqrt(14)*sqrt(77)) via torch.allclose.

Mapping: Problem+Fix sentences ⇒ $C_{\text{clear}}$; "leaving higher-rank byte-identical" ⇒ $C_{\text{correct}}$; Justification paragraph ⇒ $C_{\text{justified}}$; Tests ⇒ $C_{\text{tested}}$. All four = 1, so $Q_{\text{PR}}=1$.

Both $4.2$ GB and $8.1$ GB fit under $12$ GB, so both $k=1$ and $k=4$ are feasible. Among feasible options pick the one with least recomputation = fewest checkpoints = largest $k$. Larger $k$ means you store activations more often (checkpoint every $k$ layers), so you throw away fewer and recompute less, giving lower time cost. From the parent: $k=4$ adds only +36% time (245 ms) vs. $k=1$'s +128% (410 ms). Pick $k=4$: it fits the budget and is 1.67× faster than $k=1$. Rule: choose the largest $k$ that still fits.

Plan (all lifecycle steps):

1. Find Issue — search for an existing "add gradient checkpointing" request; if none, open one.
2. Claim It — comment "I'll take this" to avoid the L1.1 duplicate trap.
3. Fork Repo — your own copy to experiment safely.
4. Create Branch — e.g. feat/grad-checkpointing, isolating the change.
5. Code + Tests — implement the checkpoint-wrapping module and a correctness test.
6. Submit PR — with the four-part body from L4.1.
7. Code Review — respond to feedback.
8. Revisions — iterate.
9. Merge.

Evidence for the feature stage: a derivation showing memory drops from $O(L\cdot B\cdot H)$ to $O(\frac{L}{k}\cdot B\cdot H)$ at a $2L$ compute cost, plus an empirical benchmark table (memory & ms across $k$) proving the trade-off is real. Derivation alone is not enough (real memory ≠ $O$-model, per L2.3); benchmarks alone don't explain why — you need both.

$Q_{\text{PR}}$ with a forgotten issue link: correct ✅, tested ✅, justified ✅... but "link the issue" is what lets reviewers verify the change is wanted and understandable in context — its absence drops $C_{\text{clear}}$ to $0$. Multiplying: $Q_{\text{PR}}=1\times 0\times 1\times 1=0.$ A one-line omission still zeroes the score — add the link and you are back to $1$.

Expected number: the true cosine similarity $\frac{32}{\sqrt{14}\sqrt{77}}\approx \mathrm{0.97463.}$ Minimal assertion:

a = torch.tensor([1., 2., 3.]); b = torch.tensor([4., 5., 6.])
result = F.cosine_similarity(a, b, dim=0)
expected = 32. / (torch.sqrt(torch.tensor(14.)) * torch.sqrt(torch.tensor(77.)))
assert torch.allclose(result, expected)   # ~0.97463

• On buggy code: it never reaches the assert — the .sum(dim=1) line raises RuntimeError first. That still counts as a failing test (the run errors out), catching the crash.
• On fixed code: the assert compares to the hand-computed $0.97463$ via allclose, so it passes and proves mathematical correctness — not merely "doesn't crash". This single test covers both the shape bug and the value correctness, which is why it is minimal-yet-sufficient.

Multiplicative PR score with one factor zero equals? ::: 0 — any missing pillar fails the whole PR. Cosine of (1,2,3) and (4,5,6) to 4 dp? ::: 0.9746 Memory kept at k=4 in the pure $O$-model? ::: 25% (a 75% saving) Largest-$k$ rule for choosing checkpoints? ::: pick the biggest $k$ that still fits the memory budget (least recomputation).