torch.nn

Active contributors: albanD, jbschlosser, mikaylagawarecki

Purpose

torch.nn is PyTorch's neural network library. It has three main pieces: Module (the base class for layers and models), the catalogue of built-in layers, and the torch.nn.functional module of stateless ops. Most user PyTorch code is built on top of torch.nn.

Directory layout

Path	Contents
torch/nn/__init__.py	Public surface
torch/nn/modules/	Per-module subdirectories (Linear, Conv, RNN, Loss, Transformer, …)
torch/nn/functional.py	Stateless functions (~89K lines)
torch/nn/parameter.py	Parameter (a Tensor + requires_grad)
torch/nn/init.py	Initialization functions
torch/nn/parallel/	DataParallel, DistributedDataParallel
torch/nn/utils/	Clip-grad, weight-norm, parametrization, prune, etc.
torch/nn/attention/	SDPA, FlexAttention
torch/nn/intrinsic/	Fused modules (Conv-BN-ReLU, Linear-ReLU, …)
torch/nn/quantized/	Quantized variants (now mostly under torch/ao/nn/quantized/)
torch/nn/qat/	QAT variants
torch/nn/quantizable/	Modules that have a clear quantized form

Key abstractions

Type	File	Purpose
nn.Module	torch/nn/modules/module.py	The base class for layers/models
nn.Parameter	torch/nn/parameter.py	A Tensor that's automatically registered as a parameter
nn.Buffer	torch/nn/parameter.py	Non-learnable per-module state
nn.functional	torch/nn/functional.py	Stateless functional API

How Module works

nn.Module is more involved than it looks. It tracks:

• Submodules (self.named_children()) — assigning an nn.Module to an attribute auto-registers it.
• Parameters (self.named_parameters()) — assigning a nn.Parameter auto-registers it.
• Buffers (self.register_buffer(...)) — non-learnable state that participates in state_dict().
• Hooks — pre/forward/backward hooks for instrumentation and for things like DDP.
• Mode (train() vs. eval()) — toggles dropout/batch-norm behaviour.
• Device / dtype — .to(device, dtype) walks all parameters/buffers/submodules.

The implementation uses __setattr__ interception to keep these dictionaries in sync.

The "stateless call" path (torch.func.functional_call) takes a Module and a parameter dict and runs forward without mutating the module — used by transforms.

Layer catalogue

torch/nn/modules/ is split by topic:

• linear.py — Linear, Bilinear, Identity.
• conv.py — Conv1d/2d/3d, ConvTranspose*, LazyConv*.
• pooling.py — Max/Avg/Adaptive pooling.
• batchnorm.py, instancenorm.py, normalization.py — normalization layers.
• activation.py — ReLU, GELU, SiLU, Sigmoid, Tanh, …
• rnn.py — RNN, LSTM, GRU.
• transformer.py — Transformer, TransformerEncoder/Decoder, MultiheadAttention.
• loss.py — CrossEntropyLoss, MSELoss, BCEWithLogitsLoss, …
• sparse.py — Embedding, EmbeddingBag.
• dropout.py, padding.py, upsampling.py, flatten.py, pixelshuffle.py, fold.py, module.py, container.py, …

Most are thin wrappers around nn.functional; the wrapper holds parameters and buffers, the function does the math.

nn.functional

The single largest file in this package (~89K lines). Stateless functions corresponding to most layers: F.linear, F.conv2d, F.relu, F.cross_entropy, F.scaled_dot_product_attention, etc. Used directly when you want to apply an op without instantiating a module (or when you want different parameters across calls).

Attention

torch/nn/attention/ is a dedicated subpackage:

• flex_attention.py — FlexAttention: a higher-order op that accepts Python score_mod and block_mask callables and lowers to fused Triton via Inductor. Active contributor: drisspg.
• bias.py — bias variants (causal, ALiBi, …).
• _utils.py — common helpers.

F.scaled_dot_product_attention is the basic SDPA entry point. The chooser between math/mem_efficient/flash backends is in aten/src/ATen/native/transformers/.

Initialization

torch/nn/init.py has the canonical inits: xavier_uniform_, kaiming_normal_, orthogonal_, etc. Each modifies the tensor in place. Contributors typically wire a sensible default in the module's __init__.

Parametrizations

torch/nn/utils/parametrize.py lets you register a parametrization on a parameter — e.g., constrain a matrix to be orthogonal, or to be the symmetric part of a learnable matrix. The framework handles the gradient w.r.t. the unconstrained underlying parameter. Active contributor: lezcano.

Pruning and weight norm

torch/nn/utils/prune.py and torch/nn/utils/weight_norm.py are mature subpackages for pruning and reparameterization. Both use the parametrization mechanism.

Where to look

File	Purpose
torch/nn/modules/module.py	Module base class (~3K lines, well-commented)
torch/nn/functional.py	Stateless functions
torch/nn/parameter.py	Parameter, Buffer
torch/nn/init.py	Initializations
torch/nn/utils/parametrize.py	Parametrizations
torch/nn/attention/flex_attention.py	FlexAttention
torch/nn/parallel/distributed.py	DDP

Where to read next

• Systems / Autograd — the engine nn.Module builds on.
• Features / Distributed training — DDP/FSDP.
• Features / Quantization — nn.quantized modules.