vondp/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 code.
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.
Or open any problem directly in Google Colab โ every notebook has an  badge.
Option 1 โ Pull the pre-built image (fastest)
docker run -p 8888:8888 -e PORT=8888 ghcr.io/duoan/torchcode:latestOption 2 โ Build locally
make runOpen <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.
๐ง 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.
๐๏ธ Architecture & Adaptation โ Put it all together
โ๏ธ Training & Optimization
๐ฏ Inference & Decoding
๐ฌ Advanced โ Differentiators
โ๏ธ 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
8. Click ๐ Reset in the toolbar โ Blank slate โ practice again!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: ~12โ16 hours spread across 3โ4 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) โ
โ โโโ 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.
โ 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. You can also click the <b>๐ Reset</b> button in the notebook toolbar at any time to restore the blank template without restarting. </details>
<details> <summary><b>Can I use Google Colab instead?</b></summary> <br> Yes! Every notebook has an <b>Open in Colab</b> 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 <b>Colab</b> toolbar button inside JupyterLab. </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 โญ
โ Buy Me a Coffee
<a href="https://buymeacoffee.com/duoan" target="_blank"><img src="https://cdn.buymeacoffee.com/buttons/default-orange.png" alt="Buy Me A Coffee" height="41" width="174"></a>
<img src="./bmc_qr.png" alt="BMC QR Code" width="150" height="150">
Scan to support
</div>
