Thanks for taking the time to work through this exercise. It gives us a concrete piece of code to discuss and work on together in the next conversation. It is not intended to be a polished production project or a test of how much JAX syntax you have memorized.
Most people spend roughly one to two hours on the take-home. Please do not worry about polishing every detail; unfinished work and ideas for what you would try next are useful discussion material.
When you are ready, let us know and we will schedule a one-hour live pairing session. There is no formal write-up or presentation to prepare. Please have your code available so we can run it and talk through your approach, a few important changes, and any tradeoffs or open questions you identified.
Use whatever tooling you would normally reach for—your preferred editor, agent tools such as Claude Code or Cursor, documentation, and web search are all fair game. We want to understand how you actually work, so be prepared to discuss the code and suggestions you choose to keep.
This is a machine learning debugging exercise. You are given a short implementation of a Transformer model in Python and JAX NumPy and asked to find and fix bugs across its training pipeline so that it performs well on a simple task.
The model is intended to learn to sort sequences of four digits. For example:
3729>2379
The starter implementation contains several defects. As shipped, it may fail to
run or learn the task reliably. Your goal is to investigate the code and get the
model working as reliably as you can. Make whatever changes you believe are
appropriate, including adding any supporting code that helps you reason about its
behavior. The held-out examples in data_files/test.txt are available if useful.
The exercise gives us a chance to discuss your understanding of ML architectures, particularly the Transformer, as well as algorithms and data. We are also interested in how you get oriented in unfamiliar code, trace symptoms back to their source, validate changes, notice edge cases, and communicate tradeoffs. Reaching a perfect result is not required.
You will need Python 3.12 or newer and uv.
uv sync
uv run python train.pytrain.py trains from scratch with a fixed seed and reports training loss. The
first run is not expected to succeed without changes. You may run it as often as
you find useful and modify or add code as you work.
If setup problems are taking meaningful time, please contact the person who sent you the exercise. Environment troubleshooting is not what we are trying to assess.
The exercise assumes general familiarity with the Transformer architecture and the basics of training language models with supervision, such as supervised fine-tuning. If you would like a refresher, these are good starting points:
No deep JAX expertise is required. The implementation stays within the NumPy-like subset, so if you can read NumPy you should be able to read this code. For a warm-up or reference:
If any part of the prompt is unclear, please ask. We are happy to clarify the scope before or during the exercise.