PyTorch interview questions and the ML coding round: the complete guide

In this article

1. Start here: the learning path
2. What an ML coding round is and what it scores
3. Questions interviewers ask on PyTorch and ML coding
4. The numerical traps that are easy to miss
5. The topics in this path
6. All 36 guides in ML coding and PyTorch
7. Follow-up questions
8. Sources

Start here: the learning path

ML coding round

An ML coding round is an interview in which a candidate writes short, runnable machine learning code, such as a layer, a loss or a decoding loop, and explains it while writing.

From scratch

Writing code "from scratch" means building a routine from arrays and basic operations instead of calling the library function that already does the job, so the interviewer can see that you understand what the library does.

Tensor shape

A tensor's shape is the tuple of its dimension sizes, for example batch, heads, sequence and embedding.

Logits

Logits are the raw, unnormalised scores a model outputs before softmax turns them into probabilities.

The path below is ordered by dependency, so each step uses only what the steps before it taught. Follow it top to bottom once, then return to whichever step an interviewer made you stumble on. Across the 284 interviewer questions we collected for this round (counted 2026-10-24), 42 mention attention, 42 mention writing something from scratch, 14 mention the key-value (KV) cache, 11 mention gradients or the backward pass and 10 ask you to debug code ¹.

Text description of the diagram

• The path has seven steps ordered by dependency.
• Tensor shapes come first because every later topic is a shape problem.
• Debugging is the last step: finding the planted bug in a transformer or a training loop.

1. Tensor shapes and broadcasting. Learn to write the shape of every intermediate value as a comment.
2. A stable softmax. The smallest task that exposes numerical thinking.
3. Scaled dot-product attention. The task that appears most often in our collected questions.
4. Masks, heads and cross-attention. Attention with the real-world variations.
5. Losses and gradient descent. Cross-entropy, a training step, the learning rate.
6. Decoding and the KV cache. How a trained model produces text, step by step.
7. Debugging. Finding the planted bug in a transformer or a training loop.

1. AI engineer coding interview questions: LRU cache, async batching and retries
2. How do you implement scaled dot-product attention from scratch?
3. How do you write multi-head attention in PyTorch from scratch?
4. Gradient descent from scratch, with warmup and cosine learning-rate schedules
5. Cross-attention vs self-attention, and what the causal mask does
6. How do you implement the KV cache in a transformer decoder?
7. How do you write k-means and k-NN from scratch in NumPy, step by step?
8. What is a numerically stable softmax, and how do you implement it?
9. How do you write the CLIP loss in PyTorch?
10. Code review interview questions: what to check, feedback and AI-written code
11. How do you implement LLM sampling from logits?
12. PyTorch and NumPy idioms interviewers test: broadcasting and views, step by step
13. MLP forward and backward passes by hand with the chain rule, step by step
14. Linear and logistic regression from scratch by gradient descent, step by step
15. FlashAttention from scratch: how it works and how to code it
16. Positional encodings from scratch: sinusoidal and rotary (RoPE), step by step
17. How to build GPT from scratch, step by step in PyTorch
18. How do you implement 2D convolution from scratch?
19. Beam search explained: how it works, and when it beats greedy decoding
20. Decision trees with Gini splits and naive Bayes from scratch, step by step
21. What is FSDP? Fully sharded data parallel, explained
22. How do you build a toy LLM inference engine?
23. Evaluation metrics from scratch in Python, step by step: F1, AUC and NDCG
24. How do you implement LoRA in a coding interview?
25. How do you code a minimal agent loop and a minimal RAG pipeline?
26. How do you implement grouped query attention (GQA)?
27. The transformer debugging round, step by step
28. The 45-minute notebook round: what to do, minute by minute
29. What is the data-parsing round in a coding interview?
30. How do you implement a BPE tokenizer from scratch?
31. Debugging training loops and pipelines that fail silently, step by step
32. Domain-flavoured implementation tasks in ML coding interviews, step by step
33. GPU kernel questions in ML coding rounds: coalescing and fusion, step by step
34. What is self-attention, and how does it work step by step?
35. Why are attention scores divided by sqrt(d_k) in a transformer?
36. How would you write a KNN search using base Python?
37. How would you critique this regression model?
38. What are some differences between TensorFlow and PyTorch?

What an ML coding round is and what it scores

An ML coding round tests whether you can turn an equation into working array code under time pressure. The task is small by design: it fits on one screen and can grow, with the interviewer adding a requirement once the first version runs.

Reported practice differs by employer, so treat any single description as one data point. The questions we collected describe tasks such as implementing attention, writing a backward pass, parsing and summarising data, and debugging a training loop; they mention different lengths: three describe a one-hour round, one a 50-minute round in three parts, and one a task finished in 15 minutes ¹, so ask how long you have. Correct shapes, vectorised code instead of Python loops, a quick test on a tiny input and narration that explains why each step exists are habits worth building for any version of the round. Employers differ, so ask at the start which libraries are allowed.

NumPy and PyTorch appear together for a practical reason. NumPy shows that you know the arithmetic; PyTorch shows that you know the framework's conventions, such as which functions take logits and which take probabilities. Prepare in NumPy first, then translate.

Questions interviewers ask on PyTorch and ML coding

The questions fall into three families: implement a named layer or algorithm, debug code that is wrong, and extend working code under a new constraint. These are ten that recur, in our own words, each with a one-line answer.

1. Implement multi-head attention without nn.MultiheadAttention. Project to queries, keys and values, split the embedding into heads, run scaled dot-product attention per head, merge and project. See scaled dot-product attention below.
2. Add a KV cache to a given transformer. Append the new key and value to stored tensors along the sequence axis and attend over all of them. See the KV cache explainer.
3. Write top-k sampling from logits. Keep the k largest logits, set the rest to negative infinity, apply softmax, sample.
4. Write a stable softmax and the cross-entropy gradient. Subtract the row maximum; the gradient with respect to the logits is the probabilities minus the one-hot target. See cross-entropy loss.
5. Implement a two-dimensional (2D) convolution with padding and stride. Pad the input, slide the window by the stride, and multiply and sum at each position.
6. Compute precision and recall without scikit-learn. Count true positives, false positives and false negatives from the predictions, then divide. See the confusion matrix page.
7. Find the planted bug in a transformer. Check shapes, then masks, then the loss, then the update. See the debugging round.
8. Explain how autograd works. PyTorch records each operation in a graph as you compute, and the backward call walks that graph in reverse applying the chain rule. See gradient descent.
9. Explain view against reshape. view needs compatible memory layout and shares storage; reshape returns a view when it can and copies when it must.
10. Show the loss decreasing on a toy problem. Fix the seed, overfit one tiny batch and print the loss every few steps.

A wider list of answered questions for AI engineers is in the AI engineer question list, and the round's place in a full loop is in interview loops. For the deeper theory behind these tasks, start at deep learning and FlashAttention.

The numerical traps that are easy to miss

Two numerical traps are easy to miss; the numbers below come from code run for this guide.²

Overflow in softmax. Softmax is $e^{x_j}/{\sum }_k{e}^{x_k}$. For the three scores 1000, 1001 and 1002, the textbook version overflows and returns nan (not a number) for all three entries. Subtracting the maximum first returns 0.09, 0.2447 and 0.6652 ². The SciPy documentation says its implementation "uses shifting to avoid overflow" (SciPy v1.18.0 documentation, read October 2026), which is the same trick.

The scale in attention. The dot product of two random vectors with unit-variance entries has a spread that grows with the vector length. In the run, the unscaled dot-product spread (its standard deviation) was 4.0 at dimension 16, 7.99 at 64 and 15.97 at 256, and dividing by the square root of the dimension brought each back to 1.0 ². The original Transformer paper (Vaswani et al., 2017) introduced this scaling, and the PyTorch default scale is the same $\frac{1}{\sqrt{E}}$, where $E$ is the query's last dimension (PyTorch 2.14 documentation, read October 2026).

Text description of the diagram

• Unscaled dot-product spread was 4.0 at dimension 16, 7.99 at 64 and 15.97 at 256.
• Dividing by the square root of the dimension gave a spread of 1.0 at every dimension.
• At dimension 64 the average largest attention weight was 0.877 unscaled and 0.362 scaled.

The widening matters because a wide score makes the softmax nearly one-hot. In an illustrative simulation of eight keys at dimension 64, the average largest attention weight was 0.877 without the scale and 0.362 with it ². A one-hot row passes almost no gradient to the other keys.

The topics in this path

Implement scaled dot-product attention from scratch

Scaled dot-product attention computes, for each query, a weighted average of the values, where the weights are a softmax over the query's scaled dot products with every key. The shapes carry the whole answer: queries are batch by query length by embedding, keys match them, so the scores come out as query length by key length, and the output has the value dimension.

Write the shape beside every line, divide the scores by the square root of the key dimension, apply the mask before the softmax, and then multiply by the values. Test on a three-token input where you can work out the answer by hand, and check that every row of weights sums to one. The Transformers hub explains why attention exists.

Causal masks, encoder against decoder, and cross-attention

A causal mask stops a position from reading later positions by setting those scores to negative infinity before the softmax, so their weights become exactly zero. PyTorch documents that behaviour for is_causal=True and raises an error if you pass both a mask and that flag. An encoder lets every position see every other; a decoder uses the causal mask because it generates left to right.

Cross-attention is the same code with queries from one sequence and keys and values from another, so the two lengths can differ and the score matrix is not square. The usual trap is a mask built for a square case being applied to a rectangular one, which fails on shape or silently masks the wrong positions. Print the mask for a four-token input before you trust it.

Multi-head attention in PyTorch

Multi-head attention splits the embedding into several smaller heads, runs attention in each, and joins the results, so each head can attend to a different pattern. With embedding size 64 and 8 heads, each head has size 8. The mistake interviewers watch for is a reshape that mixes up the head and sequence dimensions. Say the reshape out loud: batch, sequence, heads, head size, then transpose the heads forward so attention runs over the sequence axis. After attention, transpose back, reshape to batch, sequence, embedding, and apply the output projection. A reliable test is that with one head the function must equal plain attention. Variants such as grouped-query attention share keys and values across heads and are asked as follow-ups.

Numerically stable softmax

This is the task shown above: subtract the row maximum, exponentiate, divide. State the axis, because softmax over the wrong dimension returns numbers that look plausible and are wrong. The follow-up is why the answer does not change, and the one-line proof is that the shift multiplies numerator and denominator by the same constant. A second follow-up is the log-softmax, which subtracts the log of the sum of exponentials instead and is what cross-entropy uses internally. FP16 (half precision) overflows sooner than 32-bit floats, which the fp16 against bf16 page explains. Test with the 1000, 1001, 1002 case from the worked numbers above ² and with a row of equal values, which must return equal probabilities.

Gradient descent from scratch and choosing a learning rate

A gradient-descent step subtracts the learning rate times the gradient from each parameter. Show it on one parameter first, then vectorise over all of them. The gradient descent page covers the maths; in this round you also show that the loss falls on a toy problem, which is what the research-paper task in our collected questions asks for.

Start with a linear fit, print the loss at steps 0, 10 and 100, and say what you would change if it rose: a smaller learning rate, a bug in the sign of the gradient, or unscaled inputs. For classification, PyTorch's cross-entropy function takes logits and class indices (or class probabilities) and ignores target value -100 by default (PyTorch cross-entropy documentation, read October 2026); see cross-entropy loss.

Implement the KV cache

A KV cache stores each processed token's attention keys and values so that generation does not recompute them for every new token. The coding task is to keep a cache of keys and values per layer, append the new token's key and value along the sequence axis, and let the new query attend over the whole cache.

The bug to avoid is concatenating on the wrong axis, which leaves the shapes valid for a moment and then breaks the next step. Check that generating five tokens with the cache gives the same output as recomputing everything without it, which is the standard correctness test. The concept is in the KV cache explainer and the serving context for a large language model (LLM) in LLM inference, where memory limits make the cache matter.

The debugging round

A debugging round hands you broken code, such as a transformer with one planted fault or a training loop that does not learn, and asks you to find it. Work in a fixed order: shapes first, then masks, then the loss, then the update.

Typical faults are a softmax over the wrong axis, a mask applied after the softmax, probabilities passed to a loss that expects logits, gradients that are never zeroed between steps, and a forgotten switch between training and evaluation mode. Narrate the hypothesis before each check, run the smallest input that shows the fault, and fix one thing at a time so you know which change worked. In our collected questions, 10 of the 284 questions ask for this skill directly ¹.

Follow-up questions

Do I have to implement algorithms from scratch without scikit-learn?

It depends on the round, so ask at the start. A from-scratch task usually bans the library function that solves the problem, while allowing the array library; computing precision and recall from the four counts, as the confusion matrix explainer shows, is a typical example.

Which parts can I use PyTorch for and which must be NumPy?

Ask the interviewer which tensor library is allowed, then state the choice aloud. If both are allowed, write in NumPy first because the arithmetic is more visible.

What happens if my code does not run?

Say what you expect it to do, check shapes on a tiny input, and read the error line by line. Narrating the debugging process shows the skill the round is testing.

Why does my softmax return NaN?

The exponentials overflowed. Subtract each row's maximum before exponentiating, as shown in the worked numbers above. The same overflow limit is lower in half precision, covered in fp16 against bf16, and a loss that takes the log of a zero probability gives the same symptom, which cross-entropy loss explains.

How do I show the loss decreasing when I implement part of a paper?

Fix the random seed, train on a tiny batch you can overfit, and print the loss every few steps. A loss that falls on one batch proves the gradient path works; the update rule is derived on the gradient descent page.

Sources

1. Attention Is All You Need (arXiv 1706.03762 v7) — arXiv (Vaswani et al.), 12 Jun 2017
2. torch.nn.functional.scaled_dot_product_attention (PyTorch 2.14 documentation) — PyTorch documentation, 24 Oct 2026
3. torch.nn.functional.cross_entropy (PyTorch 2.14 documentation) — PyTorch documentation, 24 Oct 2026
4. scipy.special.softmax (SciPy v1.18.0 documentation) — SciPy documentation, 24 Oct 2026