nn

Base class for all neural network modules.

A `Tensor` subclass that `Module` recognises as a learnable parameter.

Lightweight handle returned by every `register_*_hook` call.

Apply a learnable affine transformation to incoming data.

Pass-through layer that returns its input unchanged.

Apply a bilinear transformation to a pair of input tensors.

Linear layer whose input dimension is inferred on the first forward call.

Linear layer with a kernel-fused non-linear activation.

Applies a 1D convolution over a sequence of input signals.

Applies a 2D convolution over a batch of images or feature maps.

Applies a 3D convolution over volumetric data (e.g. video or medical scans).

Applies a 1D transposed convolution (fractionally-strided convolution).

Applies a 2D transposed convolution (fractionally-strided convolution).

Applies a 3D transposed convolution (fractionally-strided convolution).

A `Conv1d` that infers `in_channels` from the first input.

A `Conv2d` that infers `in_channels` from the first input.

A `Conv3d` that infers `in_channels` from the first input.

A `ConvTranspose1d` that infers `in_channels` from the first input.

A `ConvTranspose2d` that infers `in_channels` from the first input.

A `ConvTranspose3d` that infers `in_channels` from the first input.

Rectified Linear Unit activation function.

Leaky Rectified Linear Unit activation function.

Exponential Linear Unit activation function.

Continuously Differentiable Exponential Linear Unit activation function.

Scaled Exponential Linear Unit activation function.

Gaussian Error Linear Unit activation function.

Sigmoid Linear Unit (Swish) activation function.

Mish activation function.

Softplus activation function.

Hard Swish activation function.

Hard Sigmoid activation function.

Sigmoid (logistic) activation function.

Hyperbolic tangent activation function.

Softmax activation function.

Softmax activation applied over the channel dimension of a 4-D tensor.

Log-Softmax activation function.

Randomized Leaky ReLU activation function.

Cosine similarity between two tensors along a specified dimension.

Pairwise $\ell_p$ distance between corresponding rows of two tensors.

ReLU6 activation function.

Parametric Rectified Linear Unit activation function.

Threshold activation function.

Hardtanh activation function.

Log-Sigmoid activation function.

Softsign activation function.

Softmin activation function.

Gated Linear Unit activation function.

Hard Shrinkage activation function.

Tanhshrink activation function.

Soft Shrinkage activation function.

Layer normalization over the trailing dimensions of the input.

Root Mean Square Layer Normalization.

Group normalization over the channel dimension.

Batch normalization over a 2-D or 3-D input `(N, C)` or `(N, C, L)`.

Batch normalization over a 4-D input `(N, C, H, W)`.

Batch normalization over a 5-D input `(N, C, D, H, W)`.

Instance normalization for 3-D input `(N, C, L)`.

Instance normalization for 4-D input `(N, C, H, W)`.

Instance normalization for 5-D input `(N, C, D, H, W)`.

Local Response Normalization (LRN) across adjacent channels.

BatchNorm1d with lazy `num_features` inference.

BatchNorm2d with lazy `num_features` inference.

BatchNorm3d with lazy `num_features` inference.

InstanceNorm1d with lazy `num_features` inference.

InstanceNorm2d with lazy `num_features` inference.

InstanceNorm3d with lazy `num_features` inference.

Applies 1-D max pooling over a sequence.

Applies 2-D max pooling over a spatial feature map.

Applies 3-D max pooling over a volumetric feature map.

Applies 1-D average pooling over a sequence.

Applies 2-D average pooling over a spatial feature map.

Applies 3-D average pooling over a volumetric feature map.

Applies adaptive 1-D average pooling to produce a fixed output length.

Applies adaptive 2-D average pooling to produce a fixed output size.

Applies adaptive 3-D average pooling to produce a fixed output volume.

Applies adaptive 2-D max pooling to produce a fixed output size.

Applies adaptive 1-D max pooling to produce a fixed output length.

Applies adaptive 3-D max pooling to produce a fixed output volume.

Applies 1-D power-average (Lp-norm) pooling over a sequence.

Applies 2-D power-average (Lp-norm) pooling over a spatial feature map.

Applies 3-D power-average (Lp-norm) pooling over a volumetric input.

Computes a partial inverse of `MaxPool1d`.

Computes a partial inverse of `MaxPool2d`.

Computes a partial inverse of `MaxPool3d`.

Applies fractional max-pooling over a 2-D spatial feature map.

Applies fractional max-pooling over a 3-D volumetric feature map.

Randomly zero individual tensor elements during training (inverted dropout).

Randomly zero entire channels during training for 3-D inputs.

Randomly zero entire feature-map channels during training for 4-D inputs.

Randomly zero entire volumetric feature-map channels during training.

Alpha dropout — element-wise dropout that preserves SELU self-normalisation.

Channel-wise alpha dropout that preserves SELU self-normalisation.

Learnable dense lookup table that maps integer token indices to vectors.

Embedding lookup table with per-bag reduction.

Precomputed cos / sin tables for rotary positional embedding.

Fixed 1-D sinusoidal positional encoding (Vaswani et al., 2017).

Fixed 2-D sinusoidal positional encoding (DETR §A.4 / Carion 2020).

Sinusoidal-frequency embedding of integer timesteps + 2-layer MLP.

Multi-head scaled dot-product attention.

Long Short-Term Memory (LSTM) recurrent layer.

Multi-layer Gated Recurrent Unit (GRU) recurrent layer.

Multi-layer Elman recurrent neural network (RNN).

Single time-step Long Short-Term Memory (LSTM) cell.

Single time-step Gated Recurrent Unit (GRU) cell.

Single time-step Elman RNN cell.

Mean squared error (MSE) loss between each element of the prediction and the target.

Mean absolute error (MAE) loss between each element of the prediction and the target.

Cross-entropy loss for multi-class classification.

Negative log-likelihood loss.

Binary cross-entropy loss.

Binary cross-entropy loss that accepts raw logits.

Huber loss — a smooth interpolation between MSE and MAE.

Smooth L1 loss — a $\beta$-parameterised Huber loss.

Kullback–Leibler divergence loss.

Triplet margin loss for metric learning.

Triplet margin loss with a user-supplied distance function.

Cosine embedding loss for learning similarity/dissimilarity.

Margin ranking loss for pairwise ranking problems.

Hinge embedding loss for binary similarity learning.

Poisson negative log-likelihood loss for count data.

Gaussian negative log-likelihood loss for heteroscedastic regression.

Connectionist Temporal Classification (CTC) loss.

Multi-class hinge (SVM-style) margin loss.

Multi-label ranking margin loss.

CamelCase alias for `MultilabelMarginLoss`.

Soft-margin binary hinge loss.

Multi-label soft-margin loss (BCE with logits averaged over classes).

An ordered container of modules applied one after another in sequence.

A list-like container that registers all child modules with the parent.

A dict-like container that registers all child modules under string keys.

A list-like container that registers `Parameter` objects with the module.

A dict-like container that registers `Parameter` objects under string keys.

Flatten a contiguous range of dimensions into a single dimension.

Expand one dimension of a tensor into multiple dimensions.

Extract sliding local blocks (patches) from a batched 4-D input tensor.

Combine an array of sliding local blocks back into a batched tensor.

Pad a 3-D tensor (N, C, L) along the last dimension with a constant.

Pad a 4-D tensor (N, C, H, W) on all four spatial sides with a constant.

Pad a 5-D tensor (N, C, D, H, W) on all six faces with a constant.

Pad a 3-D tensor (N, C, L) with zeros along the sequence dimension.

Pad a 4-D tensor (N, C, H, W) with zeros on all four spatial sides.

Pad a 5-D tensor (N, C, D, H, W) with zeros on all six faces.

Pad a 3-D tensor (N, C, L) by reflecting at the sequence boundaries.

Pad a 4-D tensor (N, C, H, W) by reflecting at all four spatial edges.

Pad a 5-D tensor (N, C, D, H, W) by reflecting at the volume boundaries.

Pad a 3-D tensor (N, C, L) by replicating the edge values.

Pad a 4-D tensor (N, C, H, W) by replicating the edge pixels.

Pad a 5-D tensor (N, C, D, H, W) by replicating the boundary voxels.

Pad a 3-D tensor (N, C, L) by wrapping the sequence around itself.

Pad a 4-D tensor (N, C, H, W) with wrap-around in both spatial axes.

Pad a 5-D tensor (N, C, D, H, W) with wrap-around in all three axes.

Upsample an input tensor to a given spatial size or scale factor.

Nearest-neighbour upsampling for 4-D tensors (N, C, H, W).

Bilinear upsampling for 4-D tensors (N, C, H, W) with aligned corners.

Rearrange channel-stacked sub-pixel planes into a high-resolution output.

Reverse the sub-pixel rearrangement performed by `PixelShuffle`.

Shuffle channels across groups to enable cross-group information flow.

Single transformer encoder layer: self-attention followed by a position-wise feed-forward network, with residual connections and layer normalisation.

A stack of $N$ identical `TransformerEncoderLayer` modules.

Single transformer decoder layer: masked self-attention, cross-attention over encoder memory, and a position-wise FFN — each with residual connections and layer normalisation.

A stack of $N$ identical `TransformerDecoderLayer` modules.

Full encoder-decoder Transformer architecture.