TorchCode README

The project's README from github.com/duoan/TorchCode. As one markdown file: README.md.


title: TorchCode emoji: 🔥 colorFrom: red colorTo: yellow sdk: docker app_port: 7860 pinned: false

🔥 TorchCode

Crack the PyTorch interview.

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

An interactive coding platform, but for tensors. Self-hosted. Jupyter-based. Instant feedback.

PyTorch Jupyter Docker Python License: MIT

GitHub stars GitHub Container Registry Hugging Face Spaces Problems GPU

Star History Chart


🎯 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 code.

TorchCode gives you a structured practice environment with:

Feature
🧩 41 curated problems The most frequently asked PyTorch interview topics
⚖️ Automated judge Correctness checks, gradient verification, and timing
🎨 Instant feedback Colored pass/fail per test case, just like competitive programming
💡 Hints when stuck Nudges without full spoilers
📖 Reference solutions Study optimal implementations after your attempt
📊 Progress tracking What you've solved, best times, and attempt counts
🔄 One-click reset Toolbar button to reset any notebook back to its blank template — practice the same problem as many times as you want
Open In Colab Open in Colab Every notebook has an "Open in Colab" badge + toolbar button — run problems in Google Colab with zero setup

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 — opens a full JupyterLab environment in your browser. Nothing to install.

Or open any problem directly in Google Colab — every notebook has an Open In Colab badge.

Option 0b — Use the judge in Colab (pip)

In Google Colab, install the judge from PyPI so you can run check(...) without cloning the repo:

!pip install torch-judge

Then in a notebook cell:

from torch_judge import check, status, hint, reset_progress
status()           # list all problems and your progress
check("relu")      # run tests for the "relu" task
hint("relu")       # show a hint

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

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

If the registry image is unavailable for your platform, use Option 2 instead. This is the common path on Apple Silicon / arm64.

Option 2 — Build locally

make run

make run will try the prebuilt image first and automatically fall back to a local build when needed.

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

Option 3 — Standalone Web UI (Next.js + FastAPI)

For a modern, standalone coding experience with an integrated IDE and dual-pane layout:

  1. Start Backend (FastAPI):
    pip install -r api/requirements.txt
    python -m uvicorn api.main:app --port 8000 --reload
    
  2. Start Frontend (Next.js):
    cd web
    npm install
    npm run dev
    
  3. Open http://localhost:3000 in your browser.

TorchCode UI Preview


📋 Problem Set

Frequency: 🔥 = very likely in interviews, ⭐ = commonly asked, 💡 = emerging / differentiator

🧱 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 Freq Key Concepts
1 ReLU Open In Colab relu(x) Easy 🔥 Activation functions, element-wise ops
2 Softmax Open In Colab my_softmax(x, dim) Easy 🔥 Numerical stability, exp/log tricks
16 Cross-Entropy Loss Open In Colab cross_entropy_loss(logits, targets) Easy 🔥 Log-softmax, logsumexp trick
17 Dropout Open In Colab MyDropout (nn.Module) Easy 🔥 Train/eval mode, inverted scaling
18 Embedding Open In Colab MyEmbedding (nn.Module) Easy 🔥 Lookup table, weight[indices]
19 GELU Open In Colab my_gelu(x) Easy ⭐ Gaussian error linear unit, torch.erf
20 Kaiming Init Open In Colab kaiming_init(weight) Easy ⭐ std = sqrt(2/fan_in), variance scaling
21 Gradient Clipping Open In Colab clip_grad_norm(params, max_norm) Easy ⭐ Norm-based clipping, direction preservation
31 Gradient Accumulation Open In Colab accumulated_step(model, opt, ...) Easy 💡 Micro-batching, loss scaling
40 Linear Regression Open In Colab LinearRegression (3 methods) Medium 🔥 Normal equation, GD from scratch, nn.Linear
3 Linear Layer Open In Colab SimpleLinear (nn.Module) Medium 🔥 y = xW^T + b, Kaiming init, nn.Parameter
4 LayerNorm Open In Colab my_layer_norm(x, γ, β) Medium 🔥 Normalization, running stats, affine transform
7 BatchNorm Open In Colab my_batch_norm(x, γ, β) Medium ⭐ Batch vs layer statistics, train/eval behavior
8 RMSNorm Open In Colab rms_norm(x, weight) Medium ⭐ LLaMA-style norm, simpler than LayerNorm
15 SwiGLU MLP Open In Colab SwiGLUMLP (nn.Module) Medium ⭐ Gated FFN, SiLU(gate) * up, LLaMA/Mistral-style
22 Conv2d Open In Colab my_conv2d(x, weight, ...) Medium 🔥 Convolution, unfold, stride/padding

🧠 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 Freq Key Concepts
23 Cross-Attention Open In Colab MultiHeadCrossAttention (nn.Module) Medium ⭐ Encoder-decoder, Q from decoder, K/V from encoder
5 Scaled Dot-Product Attention Open In Colab scaled_dot_product_attention(Q, K, V) Hard 🔥 softmax(QK^T/√d_k)V, the foundation of everything
6 Multi-Head Attention Open In Colab MultiHeadAttention (nn.Module) Hard 🔥 Parallel heads, split/concat, projection matrices
9 Causal Self-Attention Open In Colab causal_attention(Q, K, V) Hard 🔥 Autoregressive masking with -inf, GPT-style
10 Grouped Query Attention Open In Colab GroupQueryAttention (nn.Module) Hard ⭐ GQA (LLaMA 2), KV sharing across heads
11 Sliding Window Attention Open In Colab sliding_window_attention(Q, K, V, w) Hard ⭐ Mistral-style local attention, O(n·w) complexity
12 Linear Attention Open In Colab linear_attention(Q, K, V) Hard 💡 Kernel trick, φ(Q)(φ(K)^TV), O(n·d²)
14 KV Cache Attention Open In Colab KVCacheAttention (nn.Module) Hard 🔥 Incremental decoding, cache K/V, prefill vs decode
24 RoPE Open In Colab apply_rope(q, k) Hard 🔥 Rotary position embedding, relative position via rotation
25 Flash Attention Open In Colab flash_attention(Q, K, V, block_size) Hard 💡 Tiled attention, online softmax, memory-efficient

🏗️ Architecture & Adaptation — Put it all together

# Problem What You'll Implement Difficulty Freq Key Concepts
26 LoRA Open In Colab LoRALinear (nn.Module) Medium ⭐ Low-rank adaptation, frozen base + BA update
27 ViT Patch Embedding <a

TorchCode (duoan/TorchCode)

Searching Database…

What is this?TorchCode is a LeetCode-style practice platform for PyTorch: it gives you 41 curated coding problems where you implement ML operators and architectures (softmax, attention, GPT-2 blocks, etc.) from scratch inside Jupyter notebooks, with an automated judge that checks correctness, gradients, and timing instantly.
Is this a library I import, or a standalone platform I run?It's a standalone self-hosted coding-practice app built on Jupyter notebooks, not a library you add to your own project — you either run it locally via Docker or use it through a hosted JupyterLab environment.
How do I get it running?Easiest is launching the Hugging Face Spaces instance for a full JupyterLab in your browser with zero install, or opening any problem notebook directly in Google Colab. To self-host, run make run, which pulls a prebuilt Docker image and falls back to a local build if needed.
Any install pitfalls I should know about?The prebuilt Docker image may not be available for every platform, which is the common case on Apple Silicon / arm64 — in that case you need to fall back to building the image locally instead of pulling it.
Who is this really aimed at?It's aimed at ML engineers preparing for interviews at top AI labs, who are expected to implement core PyTorch operations like attention and normalization layers from memory.
How many problems are there and what do they cover?41 curated problems split into three groups: Fundamentals (ReLU, Softmax, LayerNorm, Conv2d, Linear Regression), Attention Mechanisms (Multi-Head, Causal, GQA, RoPE, Flash Attention), and Architecture & Adaptation (LoRA, ViT Patch Embedding and more).
Does it give hints if I get stuck?Yes, it offers hints that nudge you without revealing the full solution, and reference solutions become available to study after your own attempt.
What exactly does the auto-grader check?The automated judge runs correctness checks, verifies gradients, and measures timing, giving colored pass/fail feedback per test case similar to a competitive-programming judge.
Can I retry a problem if I mess it up?Yes, every notebook has a one-click reset button in the toolbar that restores the problem to its blank template so you can practice it as many times as you want.
Creator
duoan
Size
0.6 MB
Stars
4.7K
Updated
CommitJul 31
Age
7mo
Commits · 90 Days