duoan/TorchCode
<div align="center">
๐ฅ 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.
    
  

</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 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-judgeThen 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 hintOption 1 โ Pull the pre-built image (fastest)
docker run -p 8888:8888 -e PORT=8888 ghcr.io/duoan/torchcode:latestIf 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 runmake 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:
- Start Backend (FastAPI):
pip install -r api/requirements.txt
python -m uvicorn api.main:app --port 8000 --reload- Start Frontend (Next.js):
cd web
npm install
npm run dev- Open <http://localhost:3000> in your browser.
๐ 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.
๐ฆ 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.
- Bump version in
torch_judge/_version.py(e.g.__version__ = "0.1.1"). - Configure PyPI Trusted Publisher (one-time):
- PyPI โ Your project torch-judge โ Publishing โ Add a new pending publisher
- Owner:
duoan, Repository:TorchCode, Workflow:pypi-publish.yml, Environment: (leave empty) - Run the workflow once (push a version bump to
masteror Actions โ Publish torch-judge to PyPI โ Run workflow); PyPI will then link the publisher. - 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
<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>
๐ค Contributors
Thanks to everyone who has contributed to TorchCode.
<!-- readme: contributors -start --> <table> <tbody> <tr> <td align="center"> <a href="https://github.com/duoan"> <img src="https://avatars.githubusercontent.com/u/2378740?v=4" width="100;" alt="duoan"/> <br /> <sub><b>duoan</b></sub> </a> </td> <td align="center"> <a href="https://github.com/Ando233"> <img src="https://avatars.githubusercontent.com/u/74404658?v=4" width="100;" alt="Ando233"/> <br /> <sub><b>Ando233</b></sub> </a> </td> <td align="center"> <a href="https://github.com/abhijitmjj"> <img src="https://avatars.githubusercontent.com/u/22732909?v=4" width="100;" alt="abhijitmjj"/> <br /> <sub><b>abhijitmjj</b></sub> </a> </td> <td align="center"> <a href="https://github.com/laitifranz"> <img src="https://avatars.githubusercontent.com/u/25352428?v=4" width="100;" alt="laitifranz"/> <br /> <sub><b>laitifranz</b></sub> </a> </td> <td align="center"> <a href="https://github.com/hrlics"> <img src="https://avatars.githubusercontent.com/u/90754112?v=4" width="100;" alt="hrlics"/> <br /> <sub><b>hrlics</b></sub> </a> </td> <td align="center"> <a href="https://github.com/HareshKarnan"> <img src="https://avatars.githubusercontent.com/u/5285984?v=4" width="100;" alt="HareshKarnan"/> <br /> <sub><b>HareshKarnan</b></sub> </a> </td> </tr> <tr> <td align="center"> <a href="https://github.com/ThierryHJ"> <img src="https://avatars.githubusercontent.com/u/51846529?v=4" width="100;" alt="ThierryHJ"/> <br /> <sub><b>ThierryHJ</b></sub> </a> </td> <td align="center"> <a href="https://github.com/reidemeister94"> <img src="https://avatars.githubusercontent.com/u/28828348?v=4" width="100;" alt="reidemeister94"/> <br /> <sub><b>reidemeister94</b></sub> </a> </td> </tr> <tbody> </table> <!-- readme: contributors -end -->
Auto-generated from the GitHub contributors graph with avatars and GitHub usernames.
<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>
