pyre-code-ml-practice
Pyre Code ML Practice Platform
Skill by ara.so — Daily 2026 Skills collection.
Pyre Code is a self-hosted ML coding practice platform with 68 problems ranging from ReLU to flow matching. Users implement internals of modern AI systems (Transformers, vLLM, TRL, diffusion models) in a browser editor with instant pass/fail feedback, no GPU required.
Installation
Option A — One-liner (recommended)
git clone https://github.com/whwangovo/pyre-code.git
cd pyre-code
./setup.sh
npm run dev
setup.sh creates a .venv (prefers uv, falls back to python3 -m venv), installs all Python deps, then prints the start command.
Option B — Conda
git clone https://github.com/whwangovo/pyre-code.git
cd pyre-code
conda create -n pyre python=3.11 -y && conda activate pyre
pip install -e ".[dev]"
npm install
npm run dev
Option C — Docker
git clone https://github.com/whwangovo/pyre-code.git
cd pyre-code
docker compose up --build
Progress persists in a Docker volume. Reset with docker compose down -v.
After installation
- Grading service:
http://localhost:8000 - Web app:
http://localhost:3000
Project Structure
pyre/
├── web/ # Next.js frontend
│ ├── src/app/ # Pages and API routes
│ ├── src/components/ # UI components
│ └── src/lib/ # Utilities, problem data
├── grading_service/ # FastAPI backend (grading API)
├── torch_judge/ # Judge engine — problem definitions + test runner
│ ├── problems/ # Individual problem modules
│ └── runner.py # Test execution logic
├── setup.sh # Environment bootstrap script
├── package.json # Dev scripts (runs frontend + backend concurrently)
└── pyproject.toml # Python package config
Key Commands
# Start both frontend and backend concurrently
npm run dev
# Start only the grading service (FastAPI)
cd grading_service && uvicorn main:app --reload --port 8000
# Start only the frontend (Next.js)
cd web && npm run dev
# Run Python tests
pytest torch_judge/
# Install Python package in editable mode with dev deps
pip install -e ".[dev]"
# Docker: build and start
docker compose up --build
# Docker: stop and remove volumes (reset progress)
docker compose down -v
Configuration
Environment Variables
Create web/.env.local to override defaults:
# URL of the FastAPI grading service
GRADING_SERVICE_URL=http://localhost:8000
# SQLite database path for progress tracking
DB_PATH=./data/pyre.db
AI Help (Optional)
Copy web/.env.example to web/.env and configure:
AI_HELP_BASE_URL=https://api.openai.com/v1
AI_HELP_API_KEY=$OPENAI_API_KEY
AI_HELP_MODEL=gpt-4o-mini
Any OpenAI-compatible endpoint works: OpenAI, Anthropic via proxy, Ollama, etc. Users can also set their own key in the UI if no server-side config is present.
Problem Categories
| Category | Examples |
|---|---|
| Fundamentals | ReLU, Softmax, GELU, SwiGLU, Dropout, Embedding, Linear, Kaiming Init |
| Normalization | LayerNorm, BatchNorm, RMSNorm |
| Attention | Scaled Dot-Product, Multi-Head, Causal, GQA, Flash, Differential, MLA |
| Position Encoding | Sinusoidal PE, RoPE, ALiBi, NTK-aware RoPE |
| Architecture | GPT-2 Block, ViT Block, Conv2D, MoE, Depthwise Conv |
| Training | Adam, Cosine LR, Gradient Clipping, Mixed Precision, Activation Checkpointing |
| Distributed | Tensor Parallel, FSDP, Ring Attention |
| Inference | KV Cache, Top-k Sampling, Beam Search, Speculative Decoding, Paged Attention |
| Alignment | DPO, GRPO, PPO, Reward Model |
| Diffusion | Noise Schedule, DDIM Step, Flow Matching, adaLN-Zero |
| Adaptation | LoRA, QLoRA |
| Reasoning | MCTS, Multi-Token Prediction |
| SSM | Mamba SSM |
Adding a New Problem
Problems live in torch_judge/problems/. Each problem is a Python module with a standard structure:
# torch_judge/problems/my_new_problem.py
import torch
import torch.nn as nn
from typing import Any
PROBLEM_ID = "my_new_problem"
TITLE = "My New Problem: Implement Foo"
DIFFICULTY = "medium" # "easy" | "medium" | "hard"
CATEGORY = "Fundamentals"
DESCRIPTION = """
## My New Problem
Implement the `foo` function that does XYZ.
### Input
- `x` (Tensor): shape `(batch, dim)`
### Output
- Tensor of shape `(batch, dim)`
### Formula
$$\\text{foo}(x) = x^2 + 1$$
"""
STARTER_CODE = """
import torch
def foo(x: torch.Tensor) -> torch.Tensor:
# Your implementation here
pass
"""
REFERENCE_SOLUTION = """
import torch
def foo(x: torch.Tensor) -> torch.Tensor:
return x ** 2 + 1
"""
def make_test_cases() -> list[dict[str, Any]]:
\"\"\"Return a list of test cases, each with inputs and expected outputs.\"\"\"
cases = []
# Basic case
x = torch.tensor([[1.0, 2.0, 3.0]])
cases.append({
"input": {"x": x},
"expected": x ** 2 + 1,
"description": "Basic 1x3 tensor",
})
# Batch case
x = torch.randn(4, 16)
cases.append({
"input": {"x": x},
"expected": x ** 2 + 1,
"description": "Batch of 4, dim 16",
})
# Edge case: zeros
x = torch.zeros(2, 8)
cases.append({
"input": {"x": x},
"expected": torch.ones(2, 8),
"description": "Zero tensor",
})
return cases
def grade(submission_code: str) -> dict[str, Any]:
\"\"\"Execute submission and return grading results.\"\"\"
namespace = {}
exec(submission_code, namespace)
if "foo" not in namespace:
return {"passed": 0, "total": 0, "error": "Function 'foo' not found"}
fn = namespace["foo"]
test_cases = make_test_cases()
results = []
for i, case in enumerate(test_cases):
try:
output = fn(**case["input"])
passed = torch.allclose(output, case["expected"], atol=1e-5)
results.append({
"case": i + 1,
"description": case["description"],
"passed": passed,
"error": None if passed else f"Output mismatch: got {output}, expected {case['expected']}",
})
except Exception as e:
results.append({
"case": i + 1,
"description": case["description"],
"passed": False,
"error": str(e),
})
passed = sum(r["passed"] for r in results)
return {
"passed": passed,
"total": len(results),
"results": results,
}
Register the problem
After creating the module, register it in the problem registry (typically torch_judge/registry.py or equivalent):
from torch_judge.problems import my_new_problem
PROBLEMS = [
# ... existing problems ...
my_new_problem,
]
Grading Service API
The FastAPI grading service at http://localhost:8000 exposes:
# Health check
GET /health
# List all problems
GET /problems
# Get a specific problem
GET /problems/{problem_id}
# Submit a solution
POST /submit
Content-Type: application/json
{
"problem_id": "relu",
"code": "import torch\n\ndef relu(x):\n return torch.clamp(x, min=0)"
}
# Response
{
"problem_id": "relu",
"passed": 3,
"total": 3,
"results": [
{"case": 1, "description": "Basic positive values", "passed": true, "error": null},
{"case": 2, "description": "Negative values", "passed": true, "error": null},
{"case": 3, "description": "Mixed values", "passed": true, "error": null}
]
}
Calling the grading API from Python
import requests
response = requests.post(
"http://localhost:8000/submit",
json={
"problem_id": "softmax",
"code": """
import torch
def softmax(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
x_max = x.max(dim=dim, keepdim=True).values
x_exp = torch.exp(x - x_max)
return x_exp / x_exp.sum(dim=dim, keepdim=True)
"""
}
)
result = response.json()
print(f"Passed {result['passed']}/{result['total']} test cases")
for r in result["results"]:
status = "✓" if r["passed"] else "✗"
print(f" {status} Case {r['case']}: {r['description']}")
if r["error"]:
print(f" Error: {r['error']}")
Example Implementations
Scaled Dot-Product Attention
import torch
import torch.nn.functional as F
import math
def scaled_dot_product_attention(
q: torch.Tensor, # (batch, heads, seq, d_k)
k: torch.Tensor,
v: torch.Tensor,
mask: torch.Tensor | None = None,
) -> torch.Tensor:
d_k = q.size(-1)
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
weights = F.softmax(scores, dim=-1)
return torch.matmul(weights, v)
RMSNorm
import torch
def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
rms = x.pow(2).mean(dim=-1, keepdim=True).add(eps).sqrt()
return x / rms * weight
LoRA Linear Layer
import torch
import torch.nn as nn
class LoRALinear(nn.Module):
def __init__(self, in_features: int, out_features: int, rank: int = 4, alpha: float = 1.0):
super().__init__()
self.weight = nn.Parameter(torch.randn(out_features, in_features) * 0.02)
self.lora_A = nn.Parameter(torch.randn(rank, in_features) * 0.02)
self.lora_B = nn.Parameter(torch.zeros(out_features, rank))
self.scale = alpha / rank
def forward(self, x: torch.Tensor) -> torch.Tensor:
base = x @ self.weight.T
lora = x @ self.lora_A.T @ self.lora_B.T
return base + self.scale * lora
Cosine Learning Rate Schedule
import math
def cosine_lr(step: int, max_steps: int, lr_max: float, lr_min: float = 0.0) -> float:
if step >= max_steps:
return lr_min
progress = step / max_steps
return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * progress))
KV Cache (Inference)
import torch
from dataclasses import dataclass, field
@dataclass
class KVCache:
keys: list[torch.Tensor] = field(default_factory=list)
values: list[torch.Tensor] = field(default_factory=list)
def update(self, new_k: torch.Tensor, new_v: torch.Tensor):
self.keys.append(new_k)
self.values.append(new_v)
def get(self) -> tuple[torch.Tensor, torch.Tensor]:
return torch.cat(self.keys, dim=-2), torch.cat(self.values, dim=-2)
def __len__(self) -> int:
return len(self.keys)
Learning Paths
Choose a path based on your goal:
| Path | Focus |
|---|---|
| Transformer Internals | Activations → Normalization → Attention → GPT-2 Block |
| Attention & Position Encoding | Every attention variant + RoPE, ALiBi, NTK-RoPE |
| Train a GPT from Scratch | Embeddings → architecture → loss → optimizer → tricks |
| Inference & Distributed | KV cache, quantization, sampling, tensor parallel, FSDP |
| Alignment & Reasoning | Reward model → DPO → GRPO → PPO → MCTS |
| Vision Transformer | Conv → patch embedding → ViT block |
| Diffusion & DiT | Noise schedule → DDIM → flow matching → adaLN-Zero |
| LLM Frontier Architectures | GQA, Differential Attention, MLA, MoE, MTP |
Recommended progression:
Fundamentals → Transformer Internals → Train a GPT from Scratch
│ │
▼ ▼
Attention & PE Inference & Distributed
│ │
▼ ▼
LLM Frontier Archs Alignment & Reasoning
Troubleshooting
Grading service not reachable
# Check if the service is running
curl http://localhost:8000/health
# If not, start it manually
cd grading_service
source ../.venv/bin/activate
uvicorn main:app --reload --port 8000
Python environment issues
# Verify correct Python is active
which python && python --version # should be 3.11+
# Reinstall deps
pip install -e ".[dev]"
# With uv
uv pip install -e ".[dev]"
Frontend can't connect to grading service
Check web/.env.local:
GRADING_SERVICE_URL=http://localhost:8000
Restart Next.js after changing .env.local.
Docker: port conflicts
# Check what's on port 3000 or 8000
lsof -i :3000
lsof -i :8000
# Stop conflicting processes, then retry
docker compose up --build
Submission always fails with import errors
Ensure the submission code only uses packages available in the environment. Core deps include torch, numpy, math. Check pyproject.toml for the full list.
Progress not persisting
The SQLite DB lives at ./data/pyre.db by default. For Docker, ensure the volume is mounted:
# docker-compose.yml
volumes:
- pyre_data:/app/data
Contributing a Problem
- Create
torch_judge/problems/{problem_id}.pyusing the structure above - Include
PROBLEM_ID,TITLE,DIFFICULTY,CATEGORY,DESCRIPTION,STARTER_CODE,REFERENCE_SOLUTION,make_test_cases(), andgrade() - Register in the problem registry
- Write at least 3 test cases: basic, edge case, and a larger/random tensor case
- Verify with
pytest torch_judge/before opening a PR - Open an issue first for new categories or structural changes