Lab 01 — Custom Autograd Engine

Phase: 01 — PyTorch & TF Framework Internals | Difficulty: ⭐⭐⭐☆☆ | Time: 2–3 hours

Read ../HITCHHIKERS-GUIDE.md before starting. This lab builds intuition for how PyTorch computes gradients under the hood.

What you build

A minimal scalar-valued autograd engine that mirrors PyTorch's torch.Tensor autograd mechanism:

  • Value class wrapping scalars with +, *, tanh, relu, exp, pow operations
  • Reverse-mode automatic differentiation via .backward()
  • Dynamic computation graph construction (define-by-run)
  • Gradient accumulation across shared nodes (topological sort)

Key concepts

ConceptWhat to understand
Computation graphEach op creates a node; edges are _prev set
Chain ruledL/dx = dL/dout * dout/dx applied backwards
STEHow gradients skip through non-differentiable ops
detach()Stops gradient flow — critical for quantization

Files

FilePurpose
lab.pyStub with TODO markers — implement this
solution.pyComplete working implementation
test_lab.pypytest suite — run to verify your work
requirements.txtDependencies

Run

pip install -r requirements.txt
python lab.py      # should print results after you fill in TODOs
pytest test_lab.py -v