Reading a research codebase without drowning

Setup#

pip install torch

Verified with torch 2.5.1 on CPU. The line numbers below come from that version and will move in others — the technique is the transferable part, not the numbers.

A real codebase you already have#

You do not need to clone anything to practise this. PyTorch's Adam implementation is on your disk, it is real production research-lineage code, and it answers a question from earlier in this section: exactly where does epsilon enter the update?

import inspect
import re
import torch
import torch.optim.adam as adam_module

src = inspect.getsource(adam_module)
lines = src.splitlines()
print(f"torch {torch.__version__}")
print(f"{inspect.getsourcefile(adam_module).split('site-packages')[-1]}: {len(lines)} lines")

print("\ntop-level functions -- the map of the file:")
for name in re.findall(r"^def (\w+)", src, re.M):
    print("   ", name)

print("\nwhere does epsilon enter the update?")
for i, line in enumerate(lines, 1):
    s = line.strip()
    if s.startswith("#") or "eps" not in s:
        continue
    if "denom" in s or "add_(eps)" in s:
        print(f"  line {i}: {s}")

torch 2.5.1+cu121
\torch\optim\adam.py: 803 lines

top-level functions -- the map of the file:
    _single_tensor_adam
    _multi_tensor_adam
    _fused_adam
    adam

where does epsilon enter the update?
  line 298: eps (float, optional): term added to the denominator to improve
  line 428: denom = (max_exp_avg_sqs[i].sqrt() / bias_correction2_sqrt).add_(eps)
  line 430: denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)

Eight hundred lines, reduced to the two that matter, in under a second.

The function list is the map. Three implementations of the same algorithm — one plain, one batched across tensors, one fused into a single kernel — plus a dispatcher named adam that chooses between them. That structure is extremely common in research code: a readable reference implementation beside optimised variants. Read the readable one and ignore the rest.

Line 298 is documentation, caught by the search and discarded on sight. Lines 428 and 430 are the answer. Line 428 sits in the AMSGrad branch, using the running maximum of the second moment; line 430 is the standard path.

The answer to the question. exp_avg_sq.sqrt() / bias_correction2_sqrt is the square root of the bias-corrected second moment, and .add_(eps) puts epsilon outside it. That is Algorithm 1 of Kingma and Ba, confirming what the previous lesson claimed about where PyTorch places it. The question is now settled by evidence rather than by documentation.

The same moves at the shell#

For a repository you have cloned, the equivalents are one-liners. No output block follows this one on purpose — every line prints whatever that repository happens to contain, so the result is different for every reader.

ls                                  # README, setup, and the folder names
cat README.md | head -60            # how do they say to run it?
find . -name "*.yaml" -o -name "*.json" | grep -i conf   # where the numbers live
grep -rn "def forward" --include=*.py .     # the model
grep -rn "loss" --include=*.py . | head     # the method
python train.py --help              # every knob, with its default
git log --oneline | head -20        # what changed recently

python train.py --help is the most underrated of these. Argument parsers document defaults that no README mentions, and those defaults are frequently the hyperparameters used for the paper.

Reading git history for the "why"#

The code tells you what. Version history tells you why.

git log -S "weight_decay" finds every commit that added or removed that string. git blame path/to/file.py shows who last touched each line and in which commit, and the commit message often explains a choice that the code cannot. When a released repository has a tag or release matching the paper, check out that tag — the current main branch may be two years of unrelated evolution ahead of what produced the published numbers.

Common mistakes#

Reading top to bottom. The file order is alphabetical, not logical. Enter through the entry point and follow calls.

Trusting names. A file called utils.py frequently contains the method. A function called normalize may not normalise. Verify by running, not by reading the identifier.

Assuming the repo matches the paper. Released code often reflects a later, different version of the work. Where they disagree, note the disagreement and prefer the code for reproducing numbers — then say which you followed.

Reading without running. Print shapes at the point where data enters the model. Two printed shapes settle questions that a day of reading leaves open.

Following every import. Depth-first reading is how people drown. When a call goes somewhere uninteresting, note the name and come back only if the question demands it.

Try it yourself#

Answer a new question with the same technique: does PyTorch's SGD apply momentum before or after weight decay? Point the script at torch.optim.sgd, search for momentum and weight_decay in the update lines, and write the update as a formula. Then check your formula against the docstring at the top of the same file.

What to learn next#

• Should you adopt this new method? — deciding what to do with everything you have now learned.
• Structuring an ML repo — writing code that is cheap for the next reader.
• Forward and backward hooks — instrumenting a model you did not write.