pytorch

Version Policy

torch 2.8+ baseline (current stable: 2.12). Verify version-gated features when uncertain.

When to Load

Load in addition to python-skill and python-numerics-skill — both still apply in full: Result[T], pattern matching, ADTs, Protocol, no dynamic magic, basedpyright strict. This skill specializes those principles for the PyTorch domain, where some idioms differ (nn.Module uses inheritance, model state is mutable during training).

Philosophy

• Every tensor has a shape annotation. Float[Tensor, "batch seq dim"] not bare Tensor. A bare Tensor in a signature is a type hole — it tells the reader nothing about what flows through the network.
• jaxtyping + beartype at module boundaries. @jaxtyped(typechecker=beartype) on public API functions. Not on inner loops (too slow).
• TDD for neural networks. Write the test first: what shape should the output be? Does the loss decrease? Do gradients flow to all parameters? These tests catch bugs BEFORE you stare at a training curve for an hour.
• nn.Module is the ONE exception to "Protocol over ABC." PyTorch's module system uses inheritance by design — forward(), parameters(), state_dict() all depend on it. Accept this; enforce discipline within it.
• Immutable where PyTorch allows it. Use @torch.no_grad() for inference. Use torch.inference_mode() for tighter guarantees. Freeze parameters explicitly with param.requires_grad_(False).
• Device discipline. Never scatter .to(device) calls. Set device once, pass through typed functions. Use device: torch.device as a parameter.

Complexity Annotations (Mandatory)

Non-trivial tensor code carries a big-O comment for compute AND memory, derived DFS through the call tree: callees count toward the bound, and recursion counts activation frames as memory. Write bounds in dimension names (b·h·s²·d), not n — that is what makes them actionable for batch/sequence scaling. Trivial glue (reshapes, single elementwise ops) is exempt.

Comment doctrine

• Comments are for the reader and reviewer. Comment the public API, anything that is not obvious, and the reason behind the code written as a plain sentence. Do not add a WHY: label.
• Keep the work plan out of comments. No phases, plan steps, ticket numbers, or words like "for now", "temporary", or "later" — those go stale and someone has to find and delete them. Unfinished work lives in plan.md.
• Every function that does real work gets a header comment above it: one line saying what it is for, plus its big-O for compute AND memory (like the # O(b·h·s²·d) compute, O(b·h·s²) memory line on scaled_dot_product_attention below). Tiny obvious helpers are exempt. The big-O lives in the header, not on a line inside the body.
• Editing counts too. When you change a function, keep the comments already there, fix any your change made wrong, and add the missing header, big-O, or reason. Do this only for code you add or change, not for nearby code you did not touch.

Tensor Typing (Non-Negotiable)

from jaxtyping import Float, Int, Bool, jaxtyped
from beartype import beartype
from torch import Tensor

# Every function that takes or returns a tensor gets shape annotations
# O(b·h·s²·d) compute, O(b·h·s²) memory — materialized scores matrix dominates
@jaxtyped(typechecker=beartype)
def scaled_dot_product_attention(
    query: Float[Tensor, "batch heads seq_q dim"],
    key:   Float[Tensor, "batch heads seq_k dim"],
    value: Float[Tensor, "batch heads seq_k dim_v"],
    mask:  Bool[Tensor, "batch 1 seq_q seq_k"] | None = None,
) -> Float[Tensor, "batch heads seq_q dim_v"]:
    scale: float = query.size(-1) ** -0.5
    scores: Float[Tensor, "batch heads seq_q seq_k"] = (
        torch.matmul(query, key.transpose(-2, -1)) * scale
    )
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))
    weights: Float[Tensor, "batch heads seq_q seq_k"] = torch.softmax(scores, dim=-1)
    return torch.matmul(weights, value)

This hand-rolled version is the learning implementation. Production code uses torch.nn.functional.scaled_dot_product_attention — fused flash/memory-efficient kernels that never materialize the O(s²) scores matrix. Typed wrappers for it live in references/architectures.md.

Naming conventions for tensor dimensions:

Use consistent names project-wide. The same name in two signatures means the same runtime dimension.

nn.Module Discipline

from typing import override

import torch
import torch.nn as nn
from jaxtyping import Bool, Float, jaxtyped
from beartype import beartype
from torch import Tensor

class MultiHeadAttention(nn.Module):

    def __init__(self, d_model: int, n_heads: int) -> None:
        super().__init__()
        assert d_model % n_heads == 0, (
            f"d_model ({d_model}) not divisible by n_heads ({n_heads})"
        )
        self.d_model: int = d_model
        self.n_heads: int = n_heads
        self.head_dim: int = d_model // n_heads

        # Name parameters clearly — they appear in state_dict
        self.w_q: nn.Linear = nn.Linear(d_model, d_model, bias=False)
        self.w_k: nn.Linear = nn.Linear(d_model, d_model, bias=False)
        self.w_v: nn.Linear = nn.Linear(d_model, d_model, bias=False)
        self.w_o: nn.Linear = nn.Linear(d_model, d_model, bias=False)

    @override
    @jaxtyped(typechecker=beartype)
    def forward(
        self,
        x: Float[Tensor, "batch seq d_model"],
        mask: Bool[Tensor, "batch 1 seq seq"] | None = None,
    ) -> Float[Tensor, "batch seq d_model"]:
        batch: int = x.size(0)
        seq: int = x.size(1)

        q: Float[Tensor, "batch heads seq head_dim"] = (
            self.w_q(x).view(batch, seq, self.n_heads, self.head_dim).transpose(1, 2)
        )
        k: Float[Tensor, "batch heads seq head_dim"] = (
            self.w_k(x).view(batch, seq, self.n_heads, self.head_dim).transpose(1, 2)
        )
        v: Float[Tensor, "batch heads seq head_dim"] = (
            self.w_v(x).view(batch, seq, self.n_heads, self.head_dim).transpose(1, 2)
        )

        attn: Float[Tensor, "batch heads seq head_dim"] = (
            scaled_dot_product_attention(q, k, v, mask)
        )

        out: Float[Tensor, "batch seq d_model"] = (
            attn.transpose(1, 2).contiguous().view(batch, seq, self.d_model)
        )
        return self.w_o(out)

Module rules

1. Type every __init__ attribute. self.d_model: int = d_model.
2. super().__init__() is always the first line.
3. Type the forward signature with jaxtyping. This is the contract.
4. Use assert in __init__ for construction invariants (like divisibility).
5. Never call forward() directly. Use model(x) — runs hooks.
6. Name nn.Linear, nn.Conv2d as attributes, not in lists. Names appear in state_dict.
7. nn.ModuleList / nn.ModuleDict for dynamic collections. Plain list won't register params.
8. @override on forward (and any other overridden hook). nn.Module is the sanctioned inheritance island, so overrides there stay explicit — typing.override, enforced by reportImplicitOverride = "error" in the shipped basedpyright config.

Device Management

# Set device ONCE, pass through — cuda -> mps -> cpu keeps code portable
def pick_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")

# Move model and data at the boundary, not inside functions
device: torch.device = pick_device()
model: MyModel = MyModel(d_model=256, n_heads=8).to(device)
x: Float[Tensor, "batch seq dim"] = data.to(device)

Functions should NOT contain .to(device) — they receive already-placed tensors.

MPS (Apple silicon)

• is_available() False has two causes: torch.backends.mps.is_built() distinguishes a non-MPS build from an unsupported macOS.
• float64 hard-crashes on MPS — keep tensors float32, never .double().
• Op not implemented? Rerun with PYTORCH_ENABLE_MPS_FALLBACK=1 (CPU fallback: correct but slow, fine for learning).
• AMP, benchmarking, memory tools: see references/mps.md before writing any mixed-precision or timing code for a Mac.

torch.compile

A CUDA-first win. On MPS the Inductor Metal backend is an early prototype with known codegen bugs — never apply it by default, and if you try it, verify numerics against eager mode. @jaxtyped/beartype wrappers graph-break under compile — compiled hot paths keep the shape annotations (documentation and static value) but skip the runtime checker. Modes, graph breaks, dynamic shapes: references/torch_compile.md.

Training Loop Pattern

from dataclasses import dataclass

from project_name.core.result import Result, ok, err  # python-skill's references/result.md

# Training state as an ADT — immutable snapshot per epoch
@dataclass(frozen=True, slots=True)
class EpochResult:
    epoch: int
    train_loss: float
    val_loss: float
    val_accuracy: float
    grad_norm: float

# One optimization pass over the loader; Err("DIVERGED") the moment loss
# goes non-finite — silent NaN propagation is ML's swallowed exception,
# and every step after the first NaN is wasted compute on poisoned weights.
def train_one_epoch(
    model: nn.Module,
    loader: DataLoader[tuple[Tensor, Tensor]],
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    max_grad_norm: float = 1.0,
) -> Result[float]:
    model.train()
    total_loss: float = 0.0
    n_batches: int = 0

    for batch_x, batch_y in loader:
        x: Float[Tensor, "batch ..."] = batch_x.to(device)
        y: Int[Tensor, "batch"] = batch_y.to(device)

        logits: Float[Tensor, "batch classes"] = model(x)
        loss: Float[Tensor, ""] = nn.functional.cross_entropy(logits, y)
        if not bool(torch.isfinite(loss)):
            return err("DIVERGED", f"non-finite loss at batch {n_batches}")

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
        optimizer.step()

        total_loss += loss.item()
        n_batches += 1

    return ok(total_loss / max(n_batches, 1))

@torch.no_grad()
def evaluate(
    model: nn.Module,
    loader: DataLoader[tuple[Tensor, Tensor]],
    device: torch.device,
) -> tuple[float, float]:
    model.eval()
    total_loss: float = 0.0
    correct: int = 0
    total: int = 0

    for batch_x, batch_y in loader:
        x: Float[Tensor, "batch ..."] = batch_x.to(device)
        y: Int[Tensor, "batch"] = batch_y.to(device)

        logits: Float[Tensor, "batch classes"] = model(x)
        loss: Float[Tensor, ""] = nn.functional.cross_entropy(logits, y)

        # loss is a per-batch MEAN — weight by batch size so the division
        # by example count yields the exact per-example mean even when the
        # last batch is ragged
        total_loss += loss.item() * y.size(0)
        preds: Int[Tensor, "batch"] = logits.argmax(dim=-1)
        correct += int((preds == y).sum().item())
        total += y.size(0)

    avg_loss: float = total_loss / max(total, 1)
    accuracy: float = correct / max(total, 1)
    return avg_loss, accuracy

Training discipline

1. model.train() before training, model.eval() before evaluation. Forgetting silently breaks batchnorm and dropout.
2. @torch.no_grad() on evaluation functions. Saves memory.
3. optimizer.zero_grad() before loss.backward(). Gradients accumulate by default.
4. Gradient clipping after backward(), before step(). Log the norm.
5. loss.item() to extract scalar. Don't hold the tensor — leaks graph memory.
6. Divergence is a typed failure. The non-finite loss guard returns err("DIVERGED", ...) — the caller decides (halve the LR, restore a checkpoint, abort), instead of discovering a flat curve an hour later.

TDD for Neural Networks

Testing ML code differs from testing business logic. Neural networks have structural invariants (shapes, gradient flow, parameter counts) that are testable before you ever train. Write tests in this order:

1. Shape tests (write FIRST, before implementing)

def test_forward_output_shape() -> None:
    model: MyModel = MyModel(d_model=64, n_heads=4)
    x: Float[Tensor, "2 10 64"] = torch.randn(2, 10, 64)
    out: Tensor = model(x)
    assert out.shape == (2, 10, 64)

def test_batch_independence() -> None:
    """Changing one example should not affect another's output."""
    model: MyModel = MyModel(d_model=64, n_heads=4)
    model.eval()
    x: Float[Tensor, "3 10 64"] = torch.randn(3, 10, 64)
    out_full: Tensor = model(x)
    out_single: Tensor = model(x[1:2])
    torch.testing.assert_close(out_full[1:2], out_single, atol=1e-5, rtol=1e-5)

2. Gradient flow tests (write after basic forward works)

def test_all_parameters_receive_gradients() -> None:
    model: MyModel = MyModel(d_model=64, n_heads=4)
    x: Float[Tensor, "2 10 64"] = torch.randn(2, 10, 64)
    target: Float[Tensor, "2 10 64"] = torch.randn(2, 10, 64)
    loss: Tensor = nn.functional.mse_loss(model(x), target)
    loss.backward()
    for name, param in model.named_parameters():
        assert param.grad is not None, f"{name} has no gradient"
        assert param.grad.abs().sum() > 0, f"{name} has zero gradient"

def test_gradients_are_finite() -> None:
    model: MyModel = MyModel(d_model=64, n_heads=4)
    loss: Tensor = nn.functional.mse_loss(
        model(torch.randn(2, 10, 64)), torch.randn(2, 10, 64),
    )
    loss.backward()
    for name, param in model.named_parameters():
        if param.grad is not None:
            assert torch.isfinite(param.grad).all(), f"{name} non-finite grad"

3. Training tests (write after gradients flow)

def test_loss_decreases() -> None:
    model: MyModel = MyModel(d_model=32, n_heads=2)
    optimizer: torch.optim.Adam = torch.optim.Adam(model.parameters(), lr=1e-3)
    x: Float[Tensor, "4 8 32"] = torch.randn(4, 8, 32)
    target: Float[Tensor, "4 8 32"] = torch.randn(4, 8, 32)
    losses: list[float] = []
    for _ in range(20):
        optimizer.zero_grad()
        loss: Tensor = nn.functional.mse_loss(model(x), target)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    assert losses[-1] < losses[0], f"loss did not decrease: {losses[0]:.4f} -> {losses[-1]:.4f}"

def test_overfit_single_batch() -> None:
    """Model should memorize one batch. If not, architecture is broken."""
    model: MyModel = MyModel(d_model=32, n_heads=2)
    optimizer: torch.optim.Adam = torch.optim.Adam(model.parameters(), lr=1e-3)
    x: Float[Tensor, "2 5 32"] = torch.randn(2, 5, 32)
    target: Float[Tensor, "2 5 32"] = torch.randn(2, 5, 32)
    for _ in range(200):
        optimizer.zero_grad()
        loss: Tensor = nn.functional.mse_loss(model(x), target)
        loss.backward()
        optimizer.step()
    assert loss.item() < 0.01, f"could not overfit: loss={loss.item():.4f}"

4. Reproducibility + behavioral tests

def test_deterministic_forward() -> None:
    torch.manual_seed(42)
    model: MyModel = MyModel(d_model=32, n_heads=2)
    model.eval()
    x: Float[Tensor, "1 5 32"] = torch.randn(1, 5, 32)
    torch.testing.assert_close(model(x), model(x))

def test_frozen_params_dont_update() -> None:
    model: MyModel = MyModel(d_model=32, n_heads=2)
    for p in model.w_q.parameters():
        p.requires_grad_(False)
    original: Tensor = model.w_q.weight.clone()
    optimizer: torch.optim.Adam = torch.optim.Adam(
        (p for p in model.parameters() if p.requires_grad), lr=1e-3,
    )
    loss: Tensor = nn.functional.mse_loss(
        model(torch.randn(2, 5, 32)), torch.randn(2, 5, 32),
    )
    loss.backward()
    optimizer.step()
    torch.testing.assert_close(model.w_q.weight, original)

def test_causal_mask_blocks_future() -> None:
    """Position i must not depend on position j > i."""
    model: CausalModel = CausalModel(d_model=32, n_heads=2)
    model.eval()
    x: Float[Tensor, "1 5 32"] = torch.randn(1, 5, 32)
    out_orig: Tensor = model(x)
    x_mod: Float[Tensor, "1 5 32"] = x.clone()
    x_mod[0, 4, :] = torch.randn(32)  # change last token
    out_mod: Tensor = model(x_mod)
    # positions 0-3 should be unaffected
    torch.testing.assert_close(out_orig[0, :4], out_mod[0, :4], atol=1e-5, rtol=1e-5)

TDD workflow for ML

Red → green → vacuous mutation check, at every level. Write tests in this order — each level builds on the previous:

1. Shape tests — define the API contract; run against a stub and watch it FAIL before implementing (red first)
2. Gradient flow — verify learning is physically possible
3. Training tests — loss decreasing + overfit single batch
4. Reproducibility — same seed = same result
5. Behavioral — causal masking, batch independence, equivariance

After green, prove the tests can fail: mutate the model — zero out the mask path, freeze a layer, detach a branch — and confirm test_causal_mask_blocks_future and the gradient-flow tests go red. ML tests are notoriously vacuous (test_loss_decreases passes even when a mask is ignored or a layer is dead); a test that survives the mutation proves nothing. Found a bug? The regression test lands BEFORE the fix. Deeper drills: references/testing_ml.md.

Autograd Awareness

# Inference: no_grad() saves memory
@torch.no_grad()
def predict(model: nn.Module, x: Float[Tensor, "batch ..."]) -> Tensor:
    model.eval()
    return model(x)

# Stricter: inference_mode() — no grad tracking at all
@torch.inference_mode()
def predict_strict(model: nn.Module, x: Float[Tensor, "batch ..."]) -> Tensor:
    model.eval()
    return model(x)

# Detach for value extraction
loss_value: float = loss.detach().item()

# Clone for independent copy
x_augmented: Float[Tensor, "batch ..."] = x.clone()

Checkpointing

from pathlib import Path
from typing import Any, TypedDict, cast

from project_name.core.result import Err, Ok, Result, err, ok, try_call

class Checkpoint(TypedDict):
    epoch: int
    model_state_dict: dict[str, Tensor]
    optimizer_state_dict: dict[str, Any]
    loss: float

def save_checkpoint(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    epoch: int,
    loss: float,
    path: Path,
) -> None:
    payload: Checkpoint = {
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss,
    }
    torch.save(payload, path)

# Restores model (+ optimizer) state from disk; returns the saved epoch.
# The filesystem is a boundary: missing path and corrupt/foreign payloads
# are Err, not exceptions — per python-skill's Result[T] doctrine.
def load_checkpoint(
    path: Path,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
    device: torch.device = torch.device("cpu"),
) -> Result[int]:
    if not path.exists():
        return err("NOT_FOUND", f"no checkpoint at {path}")
    checkpoint: Checkpoint
    match try_call(lambda: torch.load(path, map_location=device, weights_only=True)):
        case Ok(payload):
            # torch.load returns Any — cast ONCE at the load boundary, typed after
            checkpoint = cast(Checkpoint, payload)
        case Err() as e:
            return e
    model.load_state_dict(checkpoint["model_state_dict"])
    if optimizer is not None:
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    return ok(checkpoint["epoch"])

weights_only=True restricts the unpickler to tensors, primitive types, and containers — no arbitrary code execution. It is the default since torch 2.6; keep it explicit to document intent, and never pass weights_only=False on a file you did not produce.

Common Anti-Patterns

# BAD: bare Tensor — shape is a mystery
def forward(self, x: Tensor) -> Tensor: ...
# GOOD: shape-typed
def forward(self, x: Float[Tensor, "batch seq dim"]) -> Float[Tensor, "batch seq dim"]: ...

# BAD: .to(device) scattered inside functions
# GOOD: everything on the same device at the boundary

# BAD: forgetting model.eval() for inference
# GOOD: explicit mode switch + @torch.no_grad()

# BAD: holding tensor when you only need scalar
losses.append(loss)  # holds graph in memory
# GOOD:
losses.append(loss.item())

# BAD: in-place ops on leaf tensors that require grad
w += 1  # RuntimeError on a requires_grad leaf (plain data tensors are fine)
# GOOD:
w = w + 1

# BAD: print(f"loss: {loss}") every batch — forces GPU sync
# GOOD: periodic logging with .item()
if step % 100 == 0:
    print(f"step {step}: loss={loss.item():.4f}")

References

Cross-References

• python-skill/SKILL.md — core discipline (Result[T], ADTs, match/case, pyright strict)
• python-numerics-skill/SKILL.md — shape typing with jaxtyping, numpy patterns
• python-numerics-skill/references/jaxtyping.md — shape syntax cheatsheet
• python-skill/references/testing.md — general pytest patterns, hypothesis
• python-skill/references/tooling.md — basedpyright + ruff configuration