I applied online. The process took 1 week.

Shopify MLE Internship — AI-Assisted Pair Programming Interview

Format: Technical screen conducted as live pair-programming with an AI coding assistant. Shopify sent a GitHub repo with a skeleton problem.py (uv/pytest setup, 20-minute target) implementing single-head scaled dot-product self-attention in NumPy only, no masking. All tests initially failed; the task was to implement each function until the full suite passed — while using AI to generate code but staying responsible for verifying and explaining every line.

Skeleton Functions

1. softmax(x, axis=-1)

• Formula given: exp(x_i) / sum(exp(x))
• Embedded question: why subtract the max before exponentiating?
• Answer: prevents float overflow (exp() of large inputs → inf); subtracting the row max keeps exponents ≤ 0 without changing the ratio, since it cancels out.
• Had AI implement it; verified keepdims=True was used throughout so subtraction/division broadcast correctly.

2. scaled_dot_product_attention(Q, K, V)

• Formula given: softmax(QK^T / sqrt(d_k)) V
• Embedded question: what is this trying to encode?
• Answer: QK^T = similarity between queries/keys; /sqrt(d_k) prevents large-magnitude dot products from saturating softmax; softmax turns scores into a weighted-average distribution over V.
• Shape trace confirmed against provided example (batch=2, seq=3, d_k=4, d_v=4).
• Specifically checked with the AI: used swapaxes(-2,-1) instead of .T, since .T on a 3D batched array would incorrectly flip the batch axis too.

3. SimpleAttention class

• Learned projections W_q, W_k, W_v, W_o, each d_model x d_model, initialized small (*0.01).
• forward(): project to Q/K/V → attention → project through W_o.
• Embedded question: why the final W_o projection?
• Answer: lets the model re-transform/recombine the attended output rather than returning a raw weighted blend of values — becomes essential for merging heads in the multi-head case.

Test Suite (5 tests, all passed)

• test_softmax — sums to 1 (1D and 2D), correct argmax
• test_attention_shapes — output shape = (batch, seq_len, d_v) with mismatched d_k/d_v
• test_attention_weights — independently recomputed weights sum to 1
• test_self_attention_layer — shape preserved, output ≠ input (catches no-op bugs)
• test_numerical_stability — large-magnitude inputs (Q,K scaled ×10, d_k=64; softmax on [1000,1001,1002]) produce no NaN/Inf

Workflow

Went function-by-function: generate with AI → annotate tensor shapes as inline comments → run the relevant test → only proceed once I could explain why it passed, not just that it did. The interleaved conceptual questions in the docstrings meant the graded skill was directing/verifying AI output, not just producing passing code.