Top AI Labs interview prep

50 PyTorch Implementation Problems

This set was built from research into publicly reported ML-engineering interview rounds at Anthropic, OpenAI, DeepMind, and Mistral, plus the PyTorch fundamentals those rounds draw on. Problems are grouped by the lab whose round they're most representative of: the first 9 map to Anthropic-style rounds (RLHF, constitutional AI, interpretability), the next 8 to OpenAI-style rounds (sampling, decoding, RL objectives), the next 9 to DeepMind-style rounds (RL, self-play, normalization), the next 8 to Mistral-style rounds (efficient attention, MoE), and the final 16 are general cross-lab PyTorch fundamentals that show up across all of them.

How to read the grounding badges: every problem page states plainly whether its scenario is Confirmed — backed by a specific, retrieved source describing that lab's actual interview process or published technique — or General Pattern — a plausible, industry-standard PyTorch exercise that is not a confirmed detail of any one lab's actual interview loop.

What's on each problem page

  • The scenario, framed around the lab's domain (RLHF, decoding, RL, efficient attention), with its grounding badge.
  • The actual stub function signature and docstring from the corresponding .py file in ../ (the code you'd run locally to practice).
  • A step-by-step walkthrough of the solution approach — the reasoning, not just the final code.
  • A "how to recognize this pattern in general" section, for transferring the approach to a novel problem in an interview.

Running the code yourself

Every problem has a matching pair of files one level up, in ai-labs-pytorch/src/:

src/1_rlhf_kl_penalty.py            # stub — implement the function, then run it
src/1_rlhf_kl_penalty_solution.py   # reference solution — run it directly

python3 N_slug.py runs your implementation against the built-in test cases and asserts they pass. python3 N_slug_solution.py runs the reference solution and prints a step-by-step trace of its execution.

Start with Problem 1: KL-Divergence Penalty for RLHF-Style Policy Training, or jump to any problem from the sidebar.