Teddygat0r/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 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: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.
๐ฆ 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>
<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>
