PyTorch

torch.Tensor is the class constructor and creates an uninitialized tensor by default, inheriting the default dtype. torch.tensor (lowercase) is a factory function that infers dtype from the input data, copies the data, and is the recommended way to create tensors from Python scalars, lists, or NumPy arrays. torch.tensor always performs a data copy, while torch.as_tensor avoids the copy when possible for compatible array-protocol objects sharing memory.

t1 = torch.Tensor([1, 2])   # uses default dtype (float32)
t2 = torch.tensor([1, 2])   # infers int64 from Python ints

By default, PyTorch frees the intermediate computation graph after a single backward pass to conserve memory. Setting retain_graph=True keeps the graph alive so you can call backward() multiple times on the same graph. This is needed when computing gradients for multiple loss terms separately (e.g., in GANs with discriminator and generator losses), or when implementing custom higher-order gradient computations. Overusing it causes memory leaks, so it should be used only when necessary.

loss1.backward(retain_graph=True)  # keep graph
loss2.backward()                   # free graph

torch.no_grad() is a context manager that disables gradient computation, reducing memory usage and speeding up forward passes during inference. model.eval() switches certain layers like Dropout and BatchNorm to evaluation mode (Dropout becomes a pass-through; BatchNorm uses running statistics instead of batch statistics). Both are typically used together during inference. torch.no_grad() affects the autograd engine globally; model.eval() only changes layer behavior.

model.eval()
with torch.no_grad():
    out = model(x)

torch.compile() is a JIT compilation API that uses TorchDynamo (a Python bytecode interpreter) and TorchInductor (a code-gen backend) to compile PyTorch models into optimized kernels, achieving significant speedups with minimal code changes. It captures the computation graph, applies fusion and layout optimizations, and can target CUDA, CPU, or other backends. Unlike TorchScript, it handles arbitrary Python control flow via graph breaks. Typical speedups range from 20–80% on modern GPUs.

model = torch.compile(model)  # default mode
# modes: 'default', 'reduce-overhead', 'max-autotune'

torch.cuda.amp.autocast() automatically casts operations to lower precision (float16 or bfloat16) where numerically safe, while keeping sensitive operations (loss computation, softmax, layer norm) in float32. It works at the operator level using a whitelist/blacklist. Combined with GradScaler, which scales the loss to prevent float16 underflow during backward, AMP achieves near-FP32 accuracy with significantly reduced memory and faster throughput on Tensor Core hardware.

scaler = torch.cuda.amp.GradScaler()
with torch.autocast('cuda'):
    loss = model(x)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

parameters() returns tensors that require gradients and are updated by optimizers — these are the learnable weights. buffers() returns tensors registered via register_buffer() that are part of the module state but are not learned (e.g., BatchNorm's running_mean, running_var, or positional encoding tables). Buffers are included in state_dict() and moved with .to(device), but not returned by parameters() and therefore not updated by optimizers.

self.register_buffer('running_mean', torch.zeros(num_features))

Gradient clipping caps the magnitude of gradients before the optimizer step to prevent exploding gradients, which destabilize training especially in RNNs and deep transformers. torch.nn.utils.clip_grad_norm_() rescales gradient vectors so their global L2 norm does not exceed a threshold (typically 1.0). clip_grad_value_() clips each gradient element independently. Clipping is applied after backward() but before optimizer.step().

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

num_workers spawns separate worker processes to load and preprocess data in parallel, preventing the GPU from waiting on the CPU. pin_memory=True allocates data in page-locked (pinned) host memory, enabling faster asynchronous CPU-to-GPU transfers via DMA. Together, they create a pipeline where data is prefetched into pinned memory while the GPU processes the previous batch. On systems with limited RAM, many workers can cause memory pressure, so a balance of 4–8 workers is common.

loader = DataLoader(dataset, batch_size=32,
                    num_workers=4, pin_memory=True)

TorchScript is a statically typed subset of Python that compiles PyTorch models to an intermediate representation (IR) for deployment without a Python runtime. Its two modes are: (1) torch.jit.trace, which records operations on example inputs and works well for models without data-dependent control flow; and (2) torch.jit.script, which compiles the full Python function using type inference and handles branches and loops but requires type-annotated, TorchScript-compatible code.

# Tracing
scripted = torch.jit.trace(model, example_input)
# Scripting
scripted = torch.jit.script(model)

torch.fx is a Python-to-Python transformation toolkit that symbolically traces a PyTorch model into a Graph IR of nodes, allowing programmatic inspection and transformation (quantization, pruning, fusion). Unlike TorchScript, torch.fx stays in Python — the output is still a Python nn.Module — making transformations easier to implement. However, it requires the model to be torch.fx-traceable (no arbitrary Python constructs). It is the backbone of torch.compile's graph capture in PyTorch 2.x.

from torch.fx import symbolic_trace
graph_module = symbolic_trace(model)
print(graph_module.graph)

.detach() creates a new tensor that shares storage with the original but is detached from the computation graph — it is the safe, recommended way to get a no-grad view. .data is a legacy attribute that also returns the underlying storage without gradient tracking but bypasses autograd checks entirely, which can silently corrupt gradient computations if the tensor is modified in-place. Prefer .detach() for safe gradient-free operations.

safe = tensor.detach()    # recommended
risky = tensor.data       # legacy, can corrupt grads

Hooks are callback functions registered on tensors or modules that fire at specific points during forward/backward passes. On nn.Module, the three main types are: (1) register_forward_hook — fires after the forward pass, receives input and output; (2) register_backward_hook (deprecated; prefer register_full_backward_hook) — fires during backward, receives grad_input and grad_output; (3) register_forward_pre_hook — fires before forward, can modify inputs. Hooks are useful for debugging, feature extraction, and gradient analysis.

def hook(module, inp, out):
    print(out.shape)
handle = layer.register_forward_hook(hook)

Gradient accumulation simulates large batch training by accumulating gradients over multiple smaller micro-batches before calling optimizer.step(). Instead of zeroing gradients each step, you accumulate for N steps, then step and zero. This is essential when GPU memory cannot fit the desired batch size. Effective batch size = micro_batch_size × accumulation_steps. Care must be taken with BatchNorm statistics and learning rate scaling.

for i, (x, y) in enumerate(loader):
    loss = model(x, y) / accum_steps
    loss.backward()
    if (i+1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

A tensor is contiguous if its elements are stored in memory in row-major (C-style) order matching its stride pattern. Operations like transpose() or permute() return views with modified strides, making the tensor non-contiguous. Some operations (e.g., view()) require contiguous memory. Calling .contiguous() copies the data into a new contiguous block. .reshape() automatically calls contiguous if needed, making it safer than .view() for arbitrary tensors.

x = torch.randn(3, 4).T   # non-contiguous after transpose
x = x.contiguous()         # now safe to use .view()

nn.Sequential defines an ordered container where the forward pass feeds each module's output into the next automatically — forward() is predefined. nn.ModuleList is an index-addressable list of modules that registers them properly with PyTorch (so parameters appear in state_dict), but does not define a forward pass — you must implement it manually. Use ModuleList when modules are used conditionally or need custom data routing.

# ModuleList — manual forward needed
self.layers = nn.ModuleList([nn.Linear(d, d) for _ in range(N)])
def forward(self, x):
    for layer in self.layers: x = layer(x)
    return x

nn.LazyLinear defers weight initialization until the first forward pass, when the input feature size is known. This avoids having to manually specify in_features when building models. On the first call, the layer materializes its parameters based on the actual input shape and transitions to a regular nn.Linear. This pattern is implemented via the LazyModuleMixin, which sets a has_uninitialized_params flag. It is useful for building flexible model architectures.

layer = nn.LazyLinear(out_features=128)
# in_features inferred on first forward
out = layer(torch.randn(2, 64))

torch.vmap (now torch.func.vmap) is a functional transformation that vectorizes a function over a batch dimension without writing explicit batch loops. It applies the function to each element in a batch by transforming the computation graph to operate on batches natively. This is particularly useful for computing per-sample gradients, Jacobians, or Hessians efficiently without the overhead of looping. It composes with other functorch transforms like grad and jacrev.

from torch.func import vmap, grad
per_sample_grads = vmap(grad(loss_fn))(params, data)

Subclass torch.autograd.Function and implement static forward() and backward() methods. forward() receives ctx and input tensors, computes the output, and can save tensors for the backward pass via ctx.save_for_backward(). backward() receives ctx and output gradients, retrieves saved tensors, and returns input gradients. This is needed when implementing operations not natively supported by autograd or for numerical gradient optimizations.

class MyOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.clamp(min=0)
    @staticmethod
    def backward(ctx, grad):
        x, = ctx.saved_tensors
        return grad * (x > 0).float()

torch.save with a state_dict serializes only the learnable parameters and buffers as an OrderedDict, which is the recommended approach — it is forward-compatible and works across PyTorch versions. torch.save on the whole model uses pickle to serialize the entire object including architecture, which is fragile if class definitions change. state_dict requires the model class to be available for re-instantiation but is far more robust for production and version control.

torch.save(model.state_dict(), 'model.pt')
model.load_state_dict(torch.load('model.pt'))

torch.profiler is a performance analysis tool that captures CPU and GPU timings, memory allocations, kernel launches, and CUDA synchronization events. It uses a schedule (warm-up, active, repeat) to collect profiles during training. The trace can be exported to TensorBoard or Chrome tracing format. Key metrics include self_cpu_time_total, cuda_time_total, and memory usage per operator. It helps identify bottlenecks like slow kernels or excessive data transfer.

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA]
) as p:
    model(x)
print(p.key_averages().table(sort_by='cuda_time_total'))

PyTorch builds computation graphs dynamically (define-by-run): the graph is constructed during each forward pass, allowing Python control flow, debugging with standard debuggers, and easy model modification. TensorFlow 1.x used static graphs (define-then-run): the full graph was compiled before execution, enabling more aggressive optimization but making debugging harder. PyTorch's approach simplifies research but historically required JIT compilation (TorchScript, torch.compile) to match the deployment efficiency of static graphs.

When set to True, cuDNN runs a benchmark at startup to select the fastest convolution algorithm for the specific input shape, GPU, and cuDNN version. This adds a one-time overhead but improves throughput for fixed input shapes (e.g., static image classification). It should NOT be enabled when input shapes vary across batches (e.g., variable-length sequences or detection with different image sizes), as it will re-benchmark each new shape, hurting performance.

torch.backends.cudnn.benchmark = True  # enable for fixed input shapes

In standard Adam, adding weight_decay is mathematically equivalent to L2 regularization (decay applied to the gradient, scaled by adaptive learning rates). AdamW decouples weight decay from the gradient update — decay is applied directly to the weights after the adaptive step, not through the gradient. This makes AdamW's regularization effect consistent regardless of the gradient magnitude, which is why AdamW is preferred for training transformers and large language models.

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)

Gradient checkpointing trades compute for memory by not storing intermediate activations during the forward pass. Instead, during backward, it re-runs the forward computation for checkpointed segments to regenerate activations. checkpoint(function, *args) wraps a module or function. This reduces activation memory from O(N) to O(sqrt(N)) for N layers, enabling training of much deeper models within GPU memory constraints at the cost of ~33% extra compute.

from torch.utils.checkpoint import checkpoint
out = checkpoint(transformer_block, x)

torch.einsum implements Einstein summation notation, a concise way to express multi-dimensional tensor contractions, transpositions, and reductions using index labels. It simplifies complex operations like batch matrix multiplication, attention score computation, and outer products into readable one-liners. It is often as fast as hand-coded operations because PyTorch dispatches to optimized BLAS routines. The notation specifies input indices and output indices separated by '->'.

# Batch matmul: (B,i,k),(B,k,j)->(B,i,j)
out = torch.einsum('bik,bkj->bij', A, B)

nn.utils.parametrize (parametrizations) lets you constrain or transform a parameter during forward by registering a Parametrization module that wraps it. Each time the parameter is accessed, the parametrization transforms the raw stored value (e.g., applying spectral normalization, orthogonality constraints, or weight clipping). register_parameter simply registers a raw tensor as a learnable parameter without any transformation. Parametrizations compose cleanly and appear in state_dict properly.

torch.nn.utils.parametrize.register_parametrization(
    layer, 'weight', Orthogonal())

Autograd only tracks operations on floating-point tensors (float16, float32, float64, bfloat16). Operations on integer tensors are not tracked because gradients are not defined for discrete values. If you perform operations that mix integer and float tensors, PyTorch will error or silently not track gradients for the integer branch. You must explicitly cast integer tensors to float when they are inputs to differentiable computations.

torch.func.grad is a functional transformation from the functorch API that returns a new function computing the gradient of a scalar-output function with respect to its first argument. Unlike .backward(), which accumulates gradients into .grad attributes in-place, torch.func.grad is stateless and composable — it returns gradient tensors directly. It supports higher-order differentiation when composed with itself and works with vmap for per-sample gradients.

from torch.func import grad
grad_fn = grad(loss_fn)  # returns gradient function
g = grad_fn(params, x, y)

torch.nn.functional (F) provides stateless functions — they take all parameters and inputs explicitly and have no internal state. torch.nn modules are stateful classes that store parameters (weights, biases) and can be used in nn.Sequential or inspected via state_dict(). For example, F.relu() is a function; nn.ReLU() is a module. Use functional forms in custom nn.Module.forward() when the operation has no learnable parameters; use module forms when parameters need to be stored.

import torch.nn.functional as F
out = F.relu(x)   # stateless; no parameters stored

PyTorch follows NumPy broadcasting: tensors are compared element-wise from the trailing dimensions; dimensions of size 1 are stretched to match. A common pitfall is accidentally broadcasting when shapes are almost compatible — e.g., adding a (3,) tensor to a (3, 1) tensor gives (3, 3), not (3,). Another pitfall is (batch, seq, 1) + (1, seq, dim) unexpectedly producing (batch, seq, dim) through double broadcasting. Always verify shapes with tensor.shape before broadcasting.

a = torch.ones(3)      # shape (3,)
b = torch.ones(3, 1)   # shape (3,1)
(a + b).shape          # (3,3) — might be unintended

SyncBatchNorm synchronizes batch statistics (mean and variance) across all processes in distributed training, computing them over the global batch rather than per-process local batch. Standard BatchNorm computes statistics only from the local GPU's data, making it statistically weaker with small per-GPU batch sizes. SyncBatchNorm is used in multi-GPU training for tasks requiring accurate batch statistics, such as segmentation or GAN training. It is wrapped via torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).

model = nn.SyncBatchNorm.convert_sync_batchnorm(model)

A memory leak occurs when tensors are unintentionally kept alive — most commonly by accumulating computational graph references (e.g., appending loss.item() instead of float(loss) to a list, or storing model outputs without .detach()). Detection: monitor GPU memory with torch.cuda.memory_allocated() or nvidia-smi. Fix: call .item() or .detach() to break graph references, del tensors explicitly when done, and use torch.cuda.empty_cache() to release cached allocator memory (though it does not fix actual leaks).

# BAD: holds entire graph
losses.append(loss)
# GOOD: scalar only
losses.append(loss.item())

torch.distributed.tensor (DTensor), introduced in PyTorch 2.0, provides a distributed tensor abstraction where a single logical tensor is sharded across devices along specified dimensions. It enables tensor parallelism — splitting model weight matrices across GPUs so each device computes a portion of the matrix multiply. DTensor handles all-reduce communications automatically behind the scenes. It underpins the torch.distributed.tensor.parallel (TP) API used in large model training alongside FSDP.

GradScaler maintains a scale factor S that multiplies the loss before backward(). This shifts float16 gradients to larger magnitudes, preventing underflow. During optimizer step, scaler.step() unscales the gradients (divides by S) and checks for inf/nan; if any are found, it skips the update (to avoid corrupted parameters) and reduces S. If updates are consistently clean, S is increased. This adaptive scaling ensures training stability without requiring manual tuning.

scaler = torch.cuda.amp.GradScaler(init_scale=2**16)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

torch.cat concatenates tensors along an existing dimension, keeping the number of dimensions unchanged. torch.stack creates a new dimension and stacks tensors along it, increasing the number of dimensions by one. For example, stacking two (3, 4) tensors with dim=0 gives (2, 3, 4); concatenating them along dim=0 gives (6, 4). Use cat when merging tensors, stack when creating a batch from a list of same-shape tensors.

a, b = torch.randn(3,4), torch.randn(3,4)
torch.cat([a,b], dim=0).shape    # (6,4)
torch.stack([a,b], dim=0).shape  # (2,3,4)

Instead of calling cudaMalloc/cudaFree for every tensor, PyTorch's caching allocator maintains a pool of freed GPU memory blocks and reuses them for new allocations. This avoids the latency of CUDA synchronization points required by raw cudaMalloc. Blocks are organized by size to quickly find a suitable chunk. torch.cuda.empty_cache() releases cached-but-unused blocks back to the CUDA driver (useful for debugging memory but not needed in normal training).

print(torch.cuda.memory_allocated())   # bytes in use
print(torch.cuda.memory_reserved())    # bytes in cache

Named tensors allow dimension names (e.g., 'batch', 'height', 'width', 'channel') to be attached to tensor dimensions. Operations check name consistency and propagate names through computations. This catches bugs like accidentally summing the wrong dimension, since PyTorch will error when names mismatch. Named tensors are experimental and not widely adopted due to performance overhead, but tools like einops provide a practical alternative for readable dimension management.

t = torch.randn(2, 3, names=['batch', 'feature'])

Learning rate schedulers adjust the optimizer's lr over training by calling scheduler.step() after each epoch or step. They read the current lr from optimizer.param_groups and update it according to a policy. CosineAnnealingLR decays lr following a cosine curve to a minimum. OneCycleLR ramps lr up to a peak then decays with annealing in a single cycle. LinearWarmupCosineAnnealingLR (common in LLM training) combines a linear warmup phase with cosine decay.

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=num_epochs)
for epoch in range(num_epochs):
    train()
    scheduler.step()

init_process_group initializes the distributed communication backend (NCCL for GPUs, Gloo for CPU) and sets up the process group with world_size and rank. It must be called before any distributed operations. NCCL is the preferred GPU backend due to highly optimized all-reduce implementations. The function reads environment variables (MASTER_ADDR, MASTER_PORT) or takes them explicitly. After initialization, collective operations like all_reduce, broadcast, and all_gather become available.

torch.distributed.init_process_group(backend='nccl')

PyTorch offers three quantization approaches: (1) Post-Training Static Quantization (PTQ) — calibrates scales on a representative dataset, then fuses and quantizes ops to INT8; (2) Post-Training Dynamic Quantization — quantizes weights statically but activations dynamically per inference; (3) Quantization-Aware Training (QAT) — inserts fake quantization nodes during training to simulate INT8 effects, yielding higher accuracy. torch.ao.quantization provides the modern API; torch.quantization is the legacy interface.

import torch.ao.quantization as Q
model.qconfig = Q.get_default_qconfig('fbgemm')
Q.prepare(model, inplace=True)
calibrate(model, data)
Q.convert(model, inplace=True)

Channels Last (NHWC) stores tensor data so that channel values for a given spatial location are contiguous in memory, as opposed to the default Channels First (NCHW). NHWC matches the memory access pattern of many convolution implementations and hardware accelerators, improving cache utilization. On modern NVIDIA GPUs (Ampere+), running convolutions in NHWC format can provide significant throughput improvements. Enable via tensor.to(memory_format=torch.channels_last).

x = x.to(memory_format=torch.channels_last)
model = model.to(memory_format=torch.channels_last)

TransformerEncoder accepts two masking arguments: src_mask, a square attention bias matrix (additive, used for causal masking), and src_key_padding_mask, a boolean or float mask of shape (batch, seq) that marks padding positions. Padding positions receive -inf in the attention logits before softmax, zeroing out their attention weights. For batched training with variable-length sequences, src_key_padding_mask is essential to prevent attending to padding tokens.

# padding_mask: True where padded
out = encoder(src, src_key_padding_mask=pad_mask)

Sampler defines how indices are drawn from a dataset for each epoch. DataLoader uses a sampler to generate batches. SequentialSampler iterates indices in order (used in evaluation). RandomSampler shuffles indices each epoch. DistributedSampler partitions the dataset across processes so each GPU sees a non-overlapping subset. WeightedRandomSampler draws samples with specified probabilities to handle class imbalance. Custom samplers enable curriculum learning or hard negative mining.

from torch.utils.data import WeightedRandomSampler
sampler = WeightedRandomSampler(weights, num_samples=len(dataset))

LayerNorm normalizes across the feature dimension (within a single sample), so statistics are computed per sample and are independent of batch size — making it ideal for transformers and RNNs. BatchNorm1d normalizes across the batch dimension, requiring a sufficient batch size for stable statistics and using running averages at inference. LayerNorm is stable even with batch_size=1, while BatchNorm degrades. BatchNorm has no analog of LayerNorm's per-sample stability for variable-length sequences.

ln = nn.LayerNorm(normalized_shape=d_model)  # per sample
bn = nn.BatchNorm1d(num_features=d_model)    # across batch

zero_grad(set_to_none=True) sets parameter .grad attributes to None instead of filling them with zeros. This is faster because it avoids a memset operation over all gradient tensors. The subsequent backward pass allocates new gradient tensors only for parameters that receive a gradient. The default changed to set_to_none=True in PyTorch 2.0. It is safe for standard training but may cause issues if code explicitly checks grad is not None before conditional operations.

optimizer.zero_grad(set_to_none=True)

In distributed training launched by torchrun, the master process (rank 0) is responsible for broadcasting initial model parameters, aggregating metrics, saving checkpoints, and managing external communication (e.g., logging to TensorBoard). All ranks participate in collective communications equally, but rank 0 typically handles I/O-heavy operations to avoid duplication. torchrun (replacing torch.distributed.launch) handles fault tolerance, elastic scaling, and environment variable setup automatically.

if torch.distributed.get_rank() == 0:
    torch.save(model.state_dict(), 'ckpt.pt')

torch.logsumexp(x, dim) computes log(sum(exp(x_i))) in a numerically stable way by subtracting the maximum value before exponentiation. This prevents overflow for large inputs and underflow for very negative inputs. It is used in cross-entropy loss, softmax computation, and log-probability normalization. The manual equivalent is x_max = x.max(); result = x_max + (x - x_max).exp().sum().log(), which is exactly what torch.logsumexp implements.

log_sum = torch.logsumexp(logits, dim=-1)

Model parallelism splits different layers of a model across GPUs — GPU 0 runs layers 1-12, GPU 1 runs layers 13-24, etc. Naively this leaves all but one GPU idle (bubble problem). Pipeline parallelism improves this by splitting the mini-batch into micro-batches that flow through the pipeline, keeping all stages busy. PyTorch implements pipeline parallelism via torch.distributed.pipeline.sync.Pipe and also through the PipelineStage API in the new torch.distributed.pipelining module.

ParameterList and ParameterDict register collections of nn.Parameter objects so they appear in module.parameters() and state_dict(). Without these containers, a plain Python list or dict of Parameter objects would not be seen by PyTorch's module system, breaking optimization and saving. ParameterList supports index-based access; ParameterDict supports named access. They are analogous to ModuleList/ModuleDict but for parameters instead of modules.

self.scales = nn.ParameterList(
    [nn.Parameter(torch.ones(d)) for _ in range(n)])

MultiheadAttention projects queries, keys, and values with learned linear layers into h heads of dimension d_k = d_model/h. For each head, it computes attention scores as QK^T / sqrt(d_k), applies optional attention masking, softmax normalizes the scores, and multiplies by V. All head outputs are concatenated and projected with an output linear layer. PyTorch 2.0 uses F.scaled_dot_product_attention internally, which dispatches to FlashAttention on compatible hardware.

attn = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
out, weights = attn(query, key, value)