PyTorch tensors and autograd, deeply

import torch
x = torch.tensor([1., 2., 3.])           # CPU, float32
x = x.to('cuda')                          # move to GPU (if available)
x = torch.randn(32, 128, device='mps')    # Apple Silicon backend

If you know numpy, you know 90% of the PyTorch tensor API — .shape, .sum(), .mean(), @, broadcasting, indexing all work the same. The two superpowers numpy doesn't have:

1. Devices: a tensor lives on a specific device (cpu, cuda:0, cuda:1, mps). Operations between tensors on different devices error out by design. Move with .to(device). Don't .to() inside a tight loop — it forces a sync. 2. Autograd: any tensor with requires_grad=True records every operation that touches it into a graph, so calling .backward() on a scalar gives you gradients for free.

Under the hood, a tensor is a header (shape, dtype, device, strides, requires_grad) pointing at a storage buffer. .view(...) creates a new header sharing the storage. .contiguous() forces a memory layout copy. .clone() is a deep copy with autograd connected; .detach().clone() cuts the graph and copies. Knowing this saves you from 'why is my model out of memory' bugs.

Common dtypes in deep learning:

• float32 (fp32): default; 4 bytes/scalar. Safe, slow, big.
• float16 (fp16): 2 bytes; ~2× speed on Tensor-Core GPUs. Limited exponent range — gradients can underflow to zero. Used with GradScaler for stability.
• bfloat16 (bf16): 2 bytes; same exponent range as fp32 but worse precision. Preferred over fp16 on Ampere/Hopper/M-series — no GradScaler needed.
• float8 (fp8): 1 byte; emerging on H100 / B200 for forward + activation storage. Not yet drop-in.
• int8 / uint8: post-training quantization for inference. 4× smaller, 2-4× faster, slight accuracy hit.
• int4 (NF4): aggressive quantization used by QLoRA. Almost-no-accuracy-loss compression of frozen base models.

Mix dtypes with care. x.to(torch.bfloat16) casts; torch.autocast('cuda', dtype=torch.bfloat16): is the idiomatic mixed-precision wrapper. Rule: weights in fp32 (or bf16), activations in fp16/bf16, gradients in fp32 for the update.

w = torch.tensor(2.0, requires_grad=True)
x = torch.tensor(3.0)
y = (w * x - 5) ** 2     # forward
y.backward()              # backward — fills in w.grad
print(w.grad)             # tensor(6.0): d/dw of (w*x - 5)^2 at w=2, x=3

The computation graph is built dynamically as the forward runs — every operation that touches a leaf tensor with requires_grad=True records a node holding (1) the inputs, (2) the operation, (3) the formula for its backward pass. .backward() walks the graph in reverse, multiplying local derivatives via the chain rule, and accumulates the result into .grad on each leaf.

Leaf tensors are the ones you created directly (typically model parameters). Non-leaf tensors are computed from leaves. Only leaves get .grad populated. Calling .backward() on a non-scalar requires you to pass a vector (the upstream gradient). For scalar losses, you can just .backward().

Three rules to burn in: 1. Gradients accumulate — call optimizer.zero_grad() before each backward, or your gradients sum across batches and your model trains on nonsense. 2. In-place ops on a tensor that's part of the graph (x.add_(1)) can break autograd. PyTorch will warn or error. The fix is the non-inplace version (x = x + 1). 3. `.detach()` returns a tensor with the same data but no autograd history. Use for targets, constants, and any case where you want to use a value without backprop seeing it.

Almost every PyTorch reduction op takes a dim argument: x.sum(dim=0) sums down rows, x.sum(dim=1) sums across columns. Negative dims count from the back — dim=-1 is the last dim, which is what you want for 'sum across features within each batch element'.

x = torch.randn(32, 10)        # (batch, features)
x.mean(dim=-1, keepdim=True)   # (32, 1) — per-sample feature mean
x.softmax(dim=-1)              # (32, 10) — distribution per sample
x.topk(k=3, dim=-1)            # top-3 per sample with indices

The shape-mismatch traceback is the single most-read error in deep learning. Always use keepdim=True when you want to broadcast back. Always print x.shape liberally during development.

The einsum operator is unbeatable for tricky cases:

# Batched outer product:
torch.einsum('bi,bj->bij', a, b)        # (B, I, J)
# Multi-head attention scores:
torch.einsum('bhid,bhjd->bhij', q, k)   # (B, H, T, T)

GPU memory is the constraint that decides most architecture and batch-size choices. A few mental models:

• Activations memory ≈ batch_size × sum(activation sizes). Doubling batch size doubles activation memory.
• Gradient memory == parameter memory for full training.
• Optimizer state: Adam keeps two moments → 2× param memory. AdamW with mixed precision often 3× param memory in practice.

Total training memory ≈ params + grads + optimizer state + activations + workspace. A 7B model in fp16: 14GB params + 14GB grads + 28GB AdamW = 56GB before activations. That's why a 7B fine-tune needs an A100/H100 or aggressive tricks (LoRA, gradient checkpointing, ZeRO sharding, CPU offload).

Gradient checkpointing: instead of storing every activation for the backward pass, store only some and recompute the rest. Trades compute for memory (~30% slower, ~50% activation memory savings). torch.utils.checkpoint.checkpoint(forward_fn, *inputs).

`torch.no_grad()` for inference: disables graph construction, halves memory roughly. Pair with model.eval().

`torch.compile(model)` (PyTorch 2.x): JIT-compiles the forward graph via TorchInductor. 30-100% speedup on most workloads with one line of code.

The data path matters as much as the model. torch.utils.data.Dataset is the interface: implement __len__ and __getitem__(idx). The DataLoader wraps it with batching, shuffling, and multi-process loading.

from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    def __init__(self, paths): self.paths = paths
    def __len__(self): return len(self.paths)
    def __getitem__(self, i):
        x = load_image(self.paths[i])
        y = parse_label(self.paths[i])
        return self.transform(x), y

loader = DataLoader(MyDataset(paths), batch_size=64, shuffle=True,
                    num_workers=4, pin_memory=True, persistent_workers=True)

Key flags:

• num_workers=N: parallel data loading on N CPU workers. Rule of thumb: min(8, os.cpu_count() // 2).
• pin_memory=True: enables async DMA transfer to GPU. Always on when training on GPU.
• persistent_workers=True: don't kill+restart workers each epoch. Critical when __init__ is slow.
• prefetch_factor=2: workers preload batches ahead of consumption.

A slow DataLoader silently turns a 'GPU-bound' job CPU-bound. If nvidia-smi shows the GPU at <50% utilization during training, your data path is the bottleneck.