Skip to content

Repository files navigation

titleTorchCode
emoji🔥
colorFromred
colorToyellow
sdkdocker
app_port7860
pinnedfalse

🔥 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.

PyTorchJupyterDockerPythonLicense: MIT

GitHub starsGitHub Container RegistryHugging Face SpacesProblemsGPU

Star History Chart


🎓 このフォークについて — W2-W5 練習トラック

PracticeLangLicenseBase

duoan/TorchCode のフォークをベースにした自作練習問題集。 PyTorch で CNN 学習の主要トピックを W2-W5 の 4 週分に整理した 29 問を全問日本語化。本家の 40 問に加え、典型的な CNN 学習レシピ(pooling / augmentation / 評価指標 / 現代 optimizer 系)に直結する 16 問を spec-driven 生成インフラ と一緒に追加。

週次マッピング

Weekテーマ問題数フォルダ1問目を Colab で開く
W2MLP / 基本分類 / 基礎 optimization8practice/W2/Open In Colab
W3正則化 / 正規化 / advanced optimization9practice/W3/Open In Colab
W4CNN 基礎 + 基本 transforms7practice/W4/Open In Colab
W5CIFAR-10 advanced レシピ5practice/W5/Open In Colab

各週フォルダの README.md学習順 で問題リスト、各 .ipynb は実装 → check("...") で自動採点(5 テスト/問、計 145 テスト)。1 問目以外を開きたい時は各週の README から任意の .ipynb の Colab badge をクリック。

使い方

Colab で(推奨・セットアップ不要):practice/W{n}/ 配下の .ipynb 右上の Colab badge をクリック → Run All → ✏️ セルに実装を書く → 最後の check("...") セルで採点。

ローカル(Docker / JupyterLab):

make run # Docker 起動# ブラウザで http://localhost:8888 → practice/W2/ に移動

メンテ・拡張手順

変更タイプごとに対応する再生成スクリプトを走らせる:

やりたいこと編集対象再生成コマンド
新規問題を追加problem_specs/{id}.py を新規作成python scripts/build.py --verify
spec 問題(#41-56)の説明/テスト/解答を修正problem_specs/{id}.pypython scripts/build.py --verify
upstream 問題(#01-40)の intro 修正templates/{file}.ipynb の cell 0python scripts/sync_solutions.py
週マッピング変更scripts/week_mapping.pypython scripts/build_weeks.py
週フォルダ完全リセット(in-progress 破棄)(上記)python scripts/build_weeks.py --reset
全 56 解答の健全性チェック(なし)python scripts/verify_all_solutions.py

ソースと生成物の対応

このリポジトリは spec-driven (16 問) と upstream-hand-written (40 問) のハイブリッド。編集禁止のファイルを直接いじると次の再生成で消える。

ソース(編集 OK)生成物(編集禁止、再生成される)
problem_specs/*.py (16 問)torch_judge/tasks/{id}.py + templates/{4,5}*.ipynb + solutions/{4,5}*_solution.ipynb
templates/0*-40_*.ipynb (40 既存) cell 0対応する solutions/*_solution.ipynb の cell 0 (intro 部分のみ、code は upstream のまま)
scripts/week_mapping.pypractice/W{n}/, practice/W{n}/README.md, practice/README.md

大きな変更後は verify_all_solutions.py で 56 解答が全 pass することを確認するのが安全。

詳細は下記 Architecture / Adding Your Own Problems も参照。

License

本家 duoan/TorchCodeMIT License で公開されている。本フォークもそれを継承し MIT License で公開する。フォーク独自の追加・改変部分も MIT で利用可能。詳細は LICENSE を参照。


以下は 本家 TorchCode の README(英語、56 問全体の解説)。フォーク独自の追加問題は #41 以降。


🎯 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
🧩40 curated problemsThe most frequently asked PyTorch interview topics
⚖️Automated judgeCorrectness checks, gradient verification, and timing
🎨Instant feedbackColored pass/fail per test case, just like competitive programming
💡Hints when stuckNudges without full spoilers
📖Reference solutionsStudy optimal implementations after your attempt
📊Progress trackingWhat you've solved, best times, and attempt counts
🔄One-click resetToolbar button to reset any notebook back to its blank template — practice the same problem as many times as you want
Open In ColabOpen in ColabEvery 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 this fork's git URL (so you get the full task set, including the additions in this fork that aren't on the upstream PyPI package):

!pip install -q --force-reinstall --no-deps git+https://github.com/alextfkd/TorchCode.git

(The notebook templates already have this install cell at the top — just Run All in Colab.)

Then in a notebook cell:

fromtorch_judgeimportcheck, status, hint, reset_progressstatus() # list all problems and your progresscheck("relu") # run tests for the "relu" taskhint("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).


📋 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.

#ProblemWhat You'll ImplementDifficultyFreqKey Concepts
1ReLUOpen In Colabrelu(x)Easy🔥Activation functions, element-wise ops
2SoftmaxOpen In Colabmy_softmax(x, dim)Easy🔥Numerical stability, exp/log tricks
16Cross-Entropy LossOpen In Colabcross_entropy_loss(logits, targets)Easy🔥Log-softmax, logsumexp trick
17DropoutOpen In ColabMyDropout (nn.Module)Easy🔥Train/eval mode, inverted scaling
18EmbeddingOpen In ColabMyEmbedding (nn.Module)Easy🔥Lookup table, weight[indices]
19GELUOpen In Colabmy_gelu(x)EasyGaussian error linear unit, torch.erf
20Kaiming InitOpen In Colabkaiming_init(weight)Easystd = sqrt(2/fan_in), variance scaling
21Gradient ClippingOpen In Colabclip_grad_norm(params, max_norm)EasyNorm-based clipping, direction preservation
31Gradient AccumulationOpen In Colabaccumulated_step(model, opt, ...)Easy💡Micro-batching, loss scaling
40Linear RegressionOpen In ColabLinearRegression (3 methods)Medium🔥Normal equation, GD from scratch, nn.Linear
3Linear LayerOpen In ColabSimpleLinear (nn.Module)Medium🔥y = xW^T + b, Kaiming init, nn.Parameter
4LayerNormOpen In Colabmy_layer_norm(x, γ, β)Medium🔥Normalization, running stats, affine transform
7BatchNormOpen In Colabmy_batch_norm(x, γ, β)MediumBatch vs layer statistics, train/eval behavior
8RMSNormOpen In Colabrms_norm(x, weight)MediumLLaMA-style norm, simpler than LayerNorm
15SwiGLU MLPOpen In ColabSwiGLUMLP (nn.Module)MediumGated FFN, SiLU(gate) * up, LLaMA/Mistral-style
22Conv2dOpen In Colabmy_conv2d(x, weight, ...)Medium🔥Convolution, unfold, stride/padding
412D Max PoolingOpen In Colabmy_max_pool2d(x, k, stride, padding)Medium🔥Unfold + amax, pad with -inf for negative inputs
492D Average PoolingOpen In Colabmy_avg_pool2d(x, k, stride, padding)EasyUnfold + mean, count_include_pad=True default
50Global Average PoolingOpen In Colabglobal_avg_pool(x)Easy🔥Mean over (H, W), ResNet/MobileNet head replacing FC
51Label Smoothing CEOpen In Colablabel_smoothing_ce(logits, targets, ε)EasySmoothed target dist, modern training recipe
52Top-k AccuracyOpen In Colabtop_k_accuracy(logits, targets, k)Easy🔥topk indices + any, ImageNet eval standard
53NLL LossOpen In Colabmy_nll_loss(log_probs, targets)EasyAdvanced indexing, CE = log_softmax + NLL

🧠 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.

#ProblemWhat You'll ImplementDifficultyFreqKey Concepts
23Cross-AttentionOpen In ColabMultiHeadCrossAttention (nn.Module)MediumEncoder-decoder, Q from decoder, K/V from encoder
5Scaled Dot-Product AttentionOpen In Colabscaled_dot_product_attention(Q, K, V)Hard🔥softmax(QK^T/√d_k)V, the foundation of everything
6Multi-Head AttentionOpen In ColabMultiHeadAttention (nn.Module)Hard🔥Parallel heads, split/concat, projection matrices
9Causal Self-AttentionOpen In Colabcausal_attention(Q, K, V)Hard🔥Autoregressive masking with -inf, GPT-style
10Grouped Query AttentionOpen In ColabGroupQueryAttention (nn.Module)HardGQA (LLaMA 2), KV sharing across heads
11Sliding Window AttentionOpen In Colabsliding_window_attention(Q, K, V, w)HardMistral-style local attention, O(n·w) complexity
12Linear AttentionOpen In Colablinear_attention(Q, K, V)Hard💡Kernel trick, φ(Q)(φ(K)^TV), O(n·d²)
14KV Cache AttentionOpen In ColabKVCacheAttention (nn.Module)Hard🔥Incremental decoding, cache K/V, prefill vs decode
24RoPEOpen In Colabapply_rope(q, k)Hard🔥Rotary position embedding, relative position via rotation
25Flash AttentionOpen In Colabflash_attention(Q, K, V, block_size)Hard💡Tiled attention, online softmax, memory-efficient

🏗️ Architecture & Adaptation — Put it all together

#ProblemWhat You'll ImplementDifficultyFreqKey Concepts
26LoRAOpen In ColabLoRALinear (nn.Module)MediumLow-rank adaptation, frozen base + BA update
27ViT Patch EmbeddingOpen In ColabPatchEmbedding (nn.Module)Medium💡Image → patches → linear projection
13GPT-2 BlockOpen In ColabGPT2Block (nn.Module)HardPre-norm, causal MHA + MLP (4x, GELU), residual connections
28Mixture of ExpertsOpen In ColabMixtureOfExperts (nn.Module)HardMixtral-style, top-k routing, expert MLPs

🎨 Data Augmentation — "Boost CIFAR-10 accuracy without changing the model"

The data side of the recipe. Together with normalization + cosine LR, these turn a baseline CNN into a competitive one.

#ProblemWhat You'll ImplementDifficultyFreqKey Concepts
42Per-Channel NormalizeOpen In Colabmy_normalize(x, mean, std)Easy🔥Channel-wise (x − μ) / σ, broadcast to (C, 1, 1)
43Random Horizontal FlipOpen In Colabrandom_horizontal_flip(x, p)Easy🔥Per-sample Bernoulli mask + torch.flip
44Random Crop with PaddingOpen In Colabrandom_crop(x, size, padding)Easy🔥F.pad + per-sample random offset slice
45Cutout / RandomErasingOpen In Colabcutout(x, size)MediumRandom rectangle zero-mask, DeVries 2017
46MixupOpen In Colabmixup(x, y, α)MediumBeta(α, α), 4-tuple (x_mix, y_a, y_b, lam) interface
47CutMixOpen In Colabcutmix(x, y, α)MediumArea-based λ recomputed after boundary clipping
48TTA (Horizontal Flip)Open In Colabtta_hflip(model, x)Easy💡Probability-space averaging, free 0.3–1% bump

⚙️ Training & Optimization

#ProblemWhat You'll ImplementDifficultyFreqKey Concepts
29Adam OptimizerOpen In ColabMyAdamMediumMomentum + RMSProp, bias correction
30Cosine LR SchedulerOpen In Colabcosine_lr_schedule(step, ...)MediumLinear warmup + cosine annealing
54SGD with MomentumOpen In ColabMySGDMomentumMedium🔥v = μ·v + g (PyTorch convention — no (1−μ) factor)
55Weight Decay (L2)Open In Colabapply_weight_decay(params, wd)Easyg += wd·p, compare with decoupled WD (#56)
56AdamWOpen In ColabMyAdamWHard🔥Decoupled WD: p *= (1 − lr·λ), Transformer default

🎯 Inference & Decoding

#ProblemWhat You'll ImplementDifficultyFreqKey Concepts
32Top-k / Top-p SamplingOpen In Colabsample_top_k_top_p(logits, ...)Medium🔥Nucleus sampling, temperature scaling
33Beam SearchOpen In Colabbeam_search(log_prob_fn, ...)Medium🔥Hypothesis expansion, pruning, eos handling
34Speculative DecodingOpen In Colabspeculative_decode(target, draft, ...)Hard💡Accept/reject, draft model acceleration

🔬 Advanced — Differentiators

#ProblemWhat You'll ImplementDifficultyFreqKey Concepts
35BPE TokenizerOpen In ColabSimpleBPEHard💡Byte-pair encoding, merge rules, subword splits
36INT8 QuantizationOpen In ColabInt8Linear (nn.Module)Hard💡Per-channel quantize, scale/zero-point, buffer vs param
37DPO LossOpen In Colabdpo_loss(chosen, rejected, ...)Hard💡Direct preference optimization, alignment training
38GRPO LossOpen In Colabgrpo_loss(logps, rewards, group_ids, eps)Hard💡Group relative policy optimization, RLAIF, within-group normalized advantages
39PPO LossOpen In Colabppo_loss(new_logps, old_logps, advantages, clip_ratio)Hard💡PPO clipped surrogate loss, policy gradient, trust region

⚙️ How It Works

Each problem has two notebooks:

FilePurpose
01_relu.ipynb✏️ Blank template — write your code here
01_relu_solution.ipynb📖 Reference solution — check when stuck

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
8. Click 🔄 Reset in the toolbar → Blank slate — practice again!

In-Notebook API

fromtorch_judgeimportcheck, hint, statuscheck("relu") # Judge your implementationhint("causal_attention") # Get a hint without full spoilerstatus() # Progress dashboard — solved / attempted / todo

📅 Suggested Study Plan

Total: ~12–16 hours spread across 3–4 weeks. Perfect for interview prep on a deadline.

WeekFocusProblemsTime
1🧱 FoundationsReLU → Softmax → CE Loss → Dropout → Embedding → GELU → Linear → LayerNorm → BatchNorm → RMSNorm → SwiGLU MLP → Conv2d2–3 hrs
2🧠 Attention Deep DiveSDPA → MHA → Cross-Attn → Causal → GQA → KV Cache → Sliding Window → RoPE → Linear Attn → Flash Attn3–4 hrs
3🏗️ Architecture + TrainingGPT-2 Block → LoRA → MoE → ViT Patch → Adam → Cosine LR → Grad Clip → Grad Accumulation → Kaiming Init3–4 hrs
4🎯 Inference + AdvancedTop-k/p Sampling → Beam Search → Speculative Decoding → BPE → INT8 Quant → DPO Loss → GRPO Loss → PPO Loss + speed run3–4 hrs

🏛️ Architecture

┌──────────────────────────────────────────┐
│ Docker / Podman Container │
│ │
│ JupyterLab (:8888) │
│ ├── templates/ (reset on each run) │
│ ├── solutions/ (reference impl) │
│ ├── torch_judge/ (auto-grading) │
│ ├── torchcode-labext (JLab plugin) │
│ │ 🔄 Reset — restore template │
│ │ 🔗 Colab — open in Colab │
│ └── 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.


📦 Publishing torch-judge to PyPI (maintainers)

The judge is published as a separate package so Colab/users can pip install torch-judge without cloning the repo.

Automatic (GitHub Action)

Pushing to master after changing the package version triggers .github/workflows/pypi-publish.yml, which builds and uploads to PyPI. No git tag is required.

  1. Bump version in torch_judge/_version.py (e.g. __version__ = "0.1.1").
  2. Configure PyPI Trusted Publisher (one-time):
    • PyPI → Your project torch-judgePublishingAdd a new pending publisher
    • Owner: duoan, Repository: TorchCode, Workflow: pypi-publish.yml, Environment: (leave empty)
    • Run the workflow once (push a version bump to master or Actions → Publish torch-judge to PyPI → Run workflow); PyPI will then link the publisher.
  3. Release: commit the version bump and git push origin master.

Alternatively, use an API token: add repository secret PYPI_API_TOKEN (value = pypi-... from PyPI) and set TWINE_USERNAME=__token__ and TWINE_PASSWORD from that secret in the workflow if you prefer not to use Trusted Publishing.

Manual

pip install build twine
python -m build
twine upload dist/*

Version is in torch_judge/_version.py; bump it before each release.


❓ FAQ

Do I need a GPU?
No. Everything runs on CPU. The problems test correctness and understanding, not throughput.
Can I keep my solutions between runs?
Blank templates reset on every make run so you practice from scratch. Save your work under a different filename if you want to keep it. You can also click the 🔄 Reset button in the notebook toolbar at any time to restore the blank template without restarting.
Can I use Google Colab instead?
Yes! Every notebook has an Open in Colab badge at the top. Click it to open the problem directly in Google Colab — no Docker or local setup needed. You can also use the Colab toolbar button inside JupyterLab.
How are solutions graded?
The judge runs your function against multiple test cases using torch.allclose for numerical correctness, verifies gradients flow properly via autograd, and checks edge cases specific to each operation.
Who is this for?
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.

🤝 Contributors

Thanks to everyone who has contributed to TorchCode.

duoan
duoan
Ando233
Ando233
ThierryHJ
ThierryHJ

Auto-generated from the GitHub contributors graph with avatars and GitHub usernames.


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

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


☕ Buy Me a Coffee

Buy Me A Coffee

BMC QR Code

Scan to support

About

🔥 LeetCode for PyTorch — practice implementing softmax, attention, GPT-2 and more from scratch with instant auto-grading. Jupyter-based, self-hosted or try online.

Resources

Stars

6 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages