Tiny Autograd
Fix an autograd engine's gradients, write its broadcast rule and speed up its backward pass.
Task
This is a small autograd engine written with plain Python lists. Every value is a matrix. A program is a list of steps that build new matrices from named inputs, and the engine works out the gradient of a 1 by 1 output with respect to every input. Some gradients are wrong, the rule that handles broadcasting was never finished, and backward is too slow for deep programs. Work through it in three stages: make the failing gradient checks pass, write the broadcast rule, then make backward fast.
Requirements
- autograd(request) takes {op, inputs, steps, output} and returns a small object. inputs maps names to matrices: non-empty lists of equal-length rows of numbers (ValueError containing matrix otherwise). A step is [name, operation, input names...] with the operations matmul, add, sub, mul (item by item), tanh, sum and mean. An unknown operation raises ValueError containing step, an undefined or reused name raises ValueError containing name, and shapes that do not fit raise ValueError containing shape. Any other op raises ValueError containing op. Do not change the arguments.
- {op: "value"} returns {value}: the matrix named by output. add, sub and mul broadcast: on each axis the two sizes match, or one of them is 1 and is stretched.
- Stage 1. {op: "grad"} returns {value, grads}: the output, which must be 1 by 1 (ValueError containing output otherwise), as a number, and for every input the gradient of the output with respect to it, in the shape of that input.
- {op: "gradcheck"} returns {passed, errors}: for every input, the largest difference between the engine's gradient and a numerical one, rounded to 4 decimals. passed is true when every difference is under 0.00001.
- Stage 2. Gradients are right when add, sub or mul stretched one of its inputs: a row, a column or a 1 by 1 matrix.
- Stage 3. backward finishes quickly on a deep program in which every layer uses the layer before it twice, such as a 30-layer residual stack. Every number in a result is rounded to 6 decimals with rounded().
How you are graded
- Decomposition 15%
- Prompting the assistant 15%
- Verification 20%
- Judgment about AI output 15%
- Correctness 25%
- Communication 10%
Scores use evidence from your session. A short or incomplete session may not provide enough evidence for every skill. See a sample scorecard
Interview guidance
This is a small autograd engine. First make the failing gradient checks pass. Then finish the stubbed broadcast rule. Then make backward fast enough for deep programs. Use the AI assistant however you like. Before you change anything, run the checks and tell me how you will know a gradient is right.
- AI is allowed and expected: your own Claude, OpenAI or Gemini key, or the free in-browser model.
- You are graded on how you drive and check the AI, not on how much you prompt.
- Run the checks; read what fails.
- A new requirement will arrive partway through.
- Too little AI use is reported as “not enough to grade”, never as a fail.
You will answer 7 debrief questions at the end
- What did you delegate to the AI, and why?
- Name one AI suggestion you rejected or changed, and why.
- Choose one block the AI helped write. Explain what each line does and which test verifies it.
- What is still risky, and what would you test or ship next?
- How did you decide that a gradient was right? What did you use as ground truth?
- Which Stage 1 failures shared a cause, and which inputs told the causes apart?
- What made backward slow, and why does your version do less work for the same answer?