| `relu(x)` |  | ๐ฅ | Activation functions, element-wise ops |
| 2 | Softmax | `my_softmax(x, dim)` |  | ๐ฅ | Numerical stability, exp/log tricks |
| 16 | Cross-Entropy Loss | `cross_entropy_loss(logits, targets)` |  | ๐ฅ | Log-softmax, logsumexp trick |
| 17 | Dropout | `MyDropout` (nn.Module) |  | ๐ฅ | Train/eval mode, inverted scaling |
| 18 | Embedding | `MyEmbedding` (nn.Module) |  | ๐ฅ | Lookup table, `weight[indices]` |
| 19 | GELU | `my_gelu(x)` |  | โญ | Gaussian error linear unit, `torch.erf` |
| 20 | Kaiming Init | `kaiming_init(weight)` |  | โญ | `std = sqrt(2/fan_in)`, variance scaling |
| 21 | Gradient Clipping | `clip_grad_norm(params, max_norm)` |  | โญ | Norm-based clipping, direction preservation |
| 31 | Gradient Accumulation | `accumulated_step(model, opt, ...)` |  | ๐ก | Micro-batching, loss scaling |
| 40 | Linear Regression | `LinearRegression` (3 methods) |  | ๐ฅ | Normal equation, GD from scratch, nn.Linear |
| 3 | Linear Layer | `SimpleLinear` (nn.Module) |  | ๐ฅ | `y = xW^T + b`, Kaiming init, `nn.Parameter` |
| 4 | LayerNorm | `my_layer_norm(x, ฮณ, ฮฒ)` |  | ๐ฅ | Normalization, running stats, affine transform |
| 7 | BatchNorm | `my_batch_norm(x, ฮณ, ฮฒ)` |  | โญ | Batch vs layer statistics, train/eval behavior |
| 8 | RMSNorm | `rms_norm(x, weight)` |  | โญ | LLaMA-style norm, simpler than LayerNorm |
| 15 | SwiGLU MLP | `SwiGLUMLP` (nn.Module) |  | โญ | Gated FFN, `SiLU(gate) * up`, LLaMA/Mistral-style |
| 22 | Conv2d | `my_conv2d(x, weight, ...)` |  | ๐ฅ | 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 | `MultiHeadCrossAttention` (nn.Module) |  | โญ | Encoder-decoder, Q from decoder, K/V from encoder |
| 5 | Scaled Dot-Product Attention | `scaled_dot_product_attention(Q, K, V)` |  | ๐ฅ | `softmax(QK^T/โd_k)V`, the foundation of everything |
| 6 | Multi-Head Attention | `MultiHeadAttention` (nn.Module) |  | ๐ฅ | Parallel heads, split/concat, projection matrices |
| 9 | Causal Self-Attention | `causal_attention(Q, K, V)` |  | ๐ฅ | Autoregressive masking with `-inf`, GPT-style |
| 10 | Grouped Query Attention | `GroupQueryAttention` (nn.Module) |  | โญ | GQA (LLaMA 2), KV sharing across heads |
| 11 | Sliding Window Attention | `sliding_window_attention(Q, K, V, w)` |  | โญ | Mistral-style local attention, O(nยทw) complexity |
| 12 | Linear Attention | `linear_attention(Q, K, V)` |  | ๐ก | Kernel trick, `ฯ(Q)(ฯ(K)^TV)`, O(nยทdยฒ) |
| 14 | KV Cache Attention | `KVCacheAttention` (nn.Module) |  | ๐ฅ | Incremental decoding, cache K/V, prefill vs decode |
| 24 | RoPE | `apply_rope(q, k)` |  | ๐ฅ | Rotary position embedding, relative position via rotation |
| 25 | Flash Attention | `flash_attention(Q, K, V, block_size)` |  | ๐ก | Tiled attention, online softmax, memory-efficient |
### ๐๏ธ Architecture & Adaptation โ Put it all together
| # | Problem | What You'll Implement | Difficulty | Freq | Key Concepts |
|:---:|---------|----------------------|:----------:|:----:|--------------|
| 26 | LoRA | `LoRALinear` (nn.Module) |  | โญ | Low-rank adaptation, frozen base + `BA` update |
| 27 | ViT Patch Embedding | `PatchEmbedding` (nn.Module) |  | ๐ก | Image โ patches โ linear projection |
| 13 | GPT-2 Block | `GPT2Block` (nn.Module) |  | โญ | Pre-norm, causal MHA + MLP (4x, GELU), residual connections |
| 28 | Mixture of Experts | `MixtureOfExperts` (nn.Module) |  | โญ | Mixtral-style, top-k routing, expert MLPs |
### โ๏ธ Training & Optimization
| # | Problem | What You'll Implement | Difficulty | Freq | Key Concepts |
|:---:|---------|----------------------|:----------:|:----:|--------------|
| 29 | Adam Optimizer | `MyAdam` |  | โญ | Momentum + RMSProp, bias correction |
| 30 | Cosine LR Scheduler | `cosine_lr_schedule(step, ...)` |  | โญ | Linear warmup + cosine annealing |
### ๐ฏ Inference & Decoding
| # | Problem | What You'll Implement | Difficulty | Freq | Key Concepts |
|:---:|---------|----------------------|:----------:|:----:|--------------|
| 32 | Top-k / Top-p Sampling | `sample_top_k_top_p(logits, ...)` |  | ๐ฅ | Nucleus sampling, temperature scaling |
| 33 | Beam Search | `beam_search(log_prob_fn, ...)` |  | ๐ฅ | Hypothesis expansion, pruning, eos handling |
| 34 | Speculative Decoding | `speculative_decode(target, draft, ...)` |  | ๐ก | Accept/reject, draft model acceleration |
### ๐ฌ Advanced โ Differentiators
| # | Problem | What You'll Implement | Difficulty | Freq | Key Concepts |
|:---:|---------|----------------------|:----------:|:----:|--------------|
| 35 | BPE Tokenizer | `SimpleBPE` |  | ๐ก | Byte-pair encoding, merge rules, subword splits |
| 36 | INT8 Quantization | `Int8Linear` (nn.Module) |  | ๐ก | Per-channel quantize, scale/zero-point, buffer vs param |
| 37 | DPO Loss | `dpo_loss(chosen, rejected, ...)` |  | ๐ก | Direct preference optimization, alignment training |
| 38 | GRPO Loss | `grpo_loss(logps, rewards, group_ids, eps)` |  | ๐ก | Group relative policy optimization, RLAIF, within-group normalized advantages |
| 39 | PPO Loss | `ppo_loss(new_logps, old_logps, advantages, clip_ratio)` |  | ๐ก | PPO clipped surrogate loss, policy gradient, trust region |
---
## โ๏ธ How It Works
Each problem has **two** notebooks:
| File | Purpose |
|------|---------|
| `01_relu.ipynb` | โ๏ธ Blank template โ write your code here |
| `01_relu_solution.ipynb` | ๐ Reference solution โ check when stuck |
### Workflow
```text
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
```python
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.**
| Week | Focus | Problems | Time |
|:----:|-------|----------|:----:|
| **1** | ๐งฑ Foundations | ReLU โ Softmax โ CE Loss โ Dropout โ Embedding โ GELU โ Linear โ LayerNorm โ BatchNorm โ RMSNorm โ SwiGLU MLP โ Conv2d | 2โ3 hrs |
| **2** | ๐ง Attention Deep Dive | SDPA โ MHA โ Cross-Attn โ Causal โ GQA โ KV Cache โ Sliding Window โ RoPE โ Linear Attn โ Flash Attn | 3โ4 hrs |
| **3** | ๐๏ธ Architecture + Training | GPT-2 Block โ LoRA โ MoE โ ViT Patch โ Adam โ Cosine LR โ Grad Clip โ Grad Accumulation โ Kaiming Init | 3โ4 hrs |
| **4** | ๐ฏ Inference + Advanced | Top-k/p Sampling โ Beam Search โ Speculative Decoding โ BPE โ INT8 Quant โ DPO Loss โ GRPO Loss โ PPO Loss + speed run | 3โ4 hrs |
---
## ๐๏ธ Architecture
```text
โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
โ 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
```bash
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/`:
```python
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`](.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-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 `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
```bash
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.
**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
*Scan to support*
*Scan to support*