YiwenC/TorchCode

<div align="center">

🔥 TorchCode

Crack the PyTorch interview.

Practice implementing operators and architectures from scratch — the exact skills top ML teams test for.

Like LeetCode, but for tensors. Self-hosted. Jupyter-based. Instant feedback.

    

 

</div>

🎯 Why TorchCode?

Top companies (Meta, Google DeepMind, OpenAI, etc.) expect ML engineers to implement core operations from memory on a whiteboard. Reading papers isn't enough — you need to write softmax, LayerNorm, MultiHeadAttention, and full Transformer blocks cold.

TorchCode gives you a structured practice environment with:

No cloud. No signup. No GPU needed. Just make run — or try it instantly on Hugging Face.

🚀 Quick Start

Option 0 — Try it online (zero install)

[Launch on Hugging Face Spaces](https://huggingface.co/spaces/duoan/TorchCode) — opens a full JupyterLab environment in your browser. Nothing to install.

Option 1 — Pull the pre-built image (fastest)

docker run -p 8888:8888 -e PORT=8888 ghcr.io/duoan/torchcode:latest

Option 2 — Build locally

make run

Open <http://localhost:8888> — that's it. Works with both Docker and Podman (auto-detected).

📋 Problem Set

🧱 Fundamentals — "Implement X from scratch"

The bread and butter of ML coding interviews. You'll be asked to write these without torch.nn.

#	Problem	What You'll Implement	Difficulty	Key Concepts
1	ReLU	relu(x)		Activation functions, element-wise ops
2	Softmax	my_softmax(x, dim)		Numerical stability, exp/log tricks
3	Linear Layer	SimpleLinear (nn.Module)		y = xW^T + b, Kaiming init, nn.Parameter
4	LayerNorm	my_layer_norm(x, γ, β)		Normalization, running stats, affine transform
7	BatchNorm	my_batch_norm(x, γ, β)		Batch vs layer statistics, train/eval behavior
8	RMSNorm	rms_norm(x, weight)		LLaMA-style norm, simpler than LayerNorm

🧠 Attention Mechanisms — The heart of modern ML interviews

If you're interviewing for any role touching LLMs or Transformers, expect at least one of these.

#	Problem	What You'll Implement	Difficulty	Key Concepts
5	Scaled Dot-Product Attention	scaled_dot_product_attention(Q, K, V)		softmax(QK^T/√d_k)V, the foundation of everything
6	Multi-Head Attention	MultiHeadAttention (nn.Module)		Parallel heads, split/concat, projection matrices
9	Causal Self-Attention	causal_attention(Q, K, V)		Autoregressive masking with -inf, GPT-style
10	Grouped Query Attention	GroupQueryAttention (nn.Module)		GQA (LLaMA 2), KV sharing across heads
11	Sliding Window Attention	sliding_window_attention(Q, K, V, w)		Mistral-style local attention, O(n·w) complexity
12	Linear Attention	linear_attention(Q, K, V)		Kernel trick, φ(Q)(φ(K)^TV), O(n·d²)

🏗️ Full Architecture — Put it all together

#	Problem	What You'll Implement	Difficulty	Key Concepts
13	GPT-2 Block	GPT2Block (nn.Module)		Pre-norm, causal MHA + MLP (4x, GELU), residual connections

⚙️ How It Works

Each problem has two notebooks:

Workflow

1. Open a blank notebook           →  Read the problem description
2. Implement your solution         →  Use only basic PyTorch ops
3. Debug freely                    →  print(x.shape), check gradients, etc.
4. Run the judge cell              →  check("relu")
5. See instant colored feedback    →  ✅ pass / ❌ fail per test case
6. Stuck? Get a nudge              →  hint("relu")
7. Review the reference solution   →  01_relu_solution.ipynb

In-Notebook API

from torch_judge import check, hint, status

check("relu")               # Judge your implementation
hint("causal_attention")    # Get a hint without full spoiler
status()                    # Progress dashboard — solved / attempted / todo

📅 Suggested Study Plan

Total: ~6–8 hours spread across 2–3 weeks. Perfect for interview prep on a deadline.

🏛️ Architecture

┌──────────────────────────────────────────┐
│           Docker / Podman Container      │
│                                          │
│  JupyterLab (:8888)                      │
│    ├── templates/  (reset on each run)   │
│    ├── solutions/  (reference impl)      │
│    ├── torch_judge/ (auto-grading)       │
│    └── PyTorch (CPU), NumPy              │
│                                          │
│  Judge checks:                           │
│    ✓ Output correctness (allclose)       │
│    ✓ Gradient flow (autograd)            │
│    ✓ Shape consistency                   │
│    ✓ Edge cases & numerical stability    │
└──────────────────────────────────────────┘

Single container. Single port. No database. No frontend framework. No GPU.

🛠️ Commands

make run    # Build & start (http://localhost:8888)
make stop   # Stop the container
make clean  # Stop + remove volumes + reset all progress

🧩 Adding Your Own Problems

TorchCode uses auto-discovery — just drop a new file in torch_judge/tasks/:

TASK = {
    "id": "my_task",
    "title": "My Custom Problem",
    "difficulty": "medium",
    "function_name": "my_function",
    "hint": "Think about broadcasting...",
    "tests": [ ... ],
}

No registration needed. The judge picks it up automatically.

❓ FAQ

<details> <summary><b>Do I need a GPU?</b></summary> <br> No. Everything runs on CPU. The problems test correctness and understanding, not throughput. </details>

<details> <summary><b>Can I keep my solutions between runs?</b></summary> <br> Blank templates reset on every <code>make run</code> so you practice from scratch. Save your work under a different filename if you want to keep it. </details>

<details> <summary><b>How are solutions graded?</b></summary> <br> The judge runs your function against multiple test cases using <code>torch.allclose</code> for numerical correctness, verifies gradients flow properly via autograd, and checks edge cases specific to each operation. </details>

<details> <summary><b>Who is this for?</b></summary> <br> Anyone preparing for ML/AI engineering interviews at top tech companies, or anyone who wants to deeply understand how PyTorch operations work under the hood. </details>

<div align="center">

Built for engineers who want to deeply understand what they build.

If this helped your interview prep, consider giving it a ⭐

</div>