Common Layers

import torch.nn as nn

nn.Linear(in_features=784, out_features=256, bias=True)
nn.Bilinear(in1_features=10, in2_features=10, out_features=5)
nn.LazyLinear(out_features=256)   # infers in_features on first forward

# 1-D (sequences, audio)
nn.Conv1d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1)

# 2-D (images)
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1, dilation=1, groups=1, bias=True)

# 3-D (video, volumetric)
nn.Conv3d(in_channels=1, out_channels=16, kernel_size=(3,3,3))

# Transposed (upsampling / decoder)
nn.ConvTranspose1d(32, 1, kernel_size=4, stride=2, padding=1)
nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1)
nn.ConvTranspose3d(16, 8, kernel_size=4, stride=2, padding=1)

# Depthwise-separable pattern (groups = in_channels)
dw = nn.Conv2d(32, 32, kernel_size=3, padding=1, groups=32)   # depthwise
pw = nn.Conv2d(32, 64, kernel_size=1)                          # pointwise

# Lazy versions (infer in_channels on first call)
nn.LazyConv2d(out_channels=64, kernel_size=3)

out = floor((in + 2*padding - dilation*(kernel-1) - 1) / stride + 1)

nn.MaxPool1d(kernel_size=2, stride=2)
nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, return_indices=False, ceil_mode=False)
nn.MaxPool3d(kernel_size=2, stride=2)

nn.AvgPool1d(kernel_size=2, stride=2)
nn.AvgPool2d(kernel_size=2, stride=2, padding=0, count_include_pad=True)
nn.AvgPool3d(kernel_size=2, stride=2)

# Adaptive — specify output size, not kernel
nn.AdaptiveMaxPool2d(output_size=(7, 7))
nn.AdaptiveAvgPool2d(output_size=(1, 1))   # global average pooling
nn.AdaptiveAvgPool1d(output_size=1)

# Fractional max pool (random pooling)
nn.FractionalMaxPool2d(kernel_size=3, output_ratio=0.5)

# Max unpooling (inverse of MaxPool with indices)
pool = nn.MaxPool2d(2, return_indices=True)
unpool = nn.MaxUnpool2d(2)
out, idx = pool(x)
recovered = unpool(out, idx)

# Batch Normalization (normalize over N,H,W — per channel)
nn.BatchNorm1d(num_features=256, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True)
nn.BatchNorm2d(num_features=64)
nn.BatchNorm3d(num_features=32)

# Layer Normalization (normalize over last D dims — per sample)
nn.LayerNorm(normalized_shape=256)
nn.LayerNorm(normalized_shape=[T, C])   # normalize over last 2 dims

# Instance Normalization (normalize over H,W — per sample per channel)
nn.InstanceNorm1d(num_features=64)
nn.InstanceNorm2d(num_features=64)

# Group Normalization (normalize over groups of channels)
nn.GroupNorm(num_groups=32, num_channels=256)
nn.GroupNorm(num_groups=1, num_channels=256)   # LayerNorm equivalent

# RMS Normalization (no mean centering — used in LLMs)
nn.RMSNorm(normalized_shape=256, eps=1e-6)

nn.RNN(input_size=10, hidden_size=32, num_layers=1, nonlinearity='tanh',
       bias=True, batch_first=False, dropout=0.0, bidirectional=False)

nn.LSTM(input_size=10, hidden_size=32, num_layers=2,
        bias=True, batch_first=True, dropout=0.2, bidirectional=False, proj_size=0)

nn.GRU(input_size=10, hidden_size=32, num_layers=1,
       batch_first=True, dropout=0.0, bidirectional=False)

# Cells (single step, no loop)
nn.RNNCell(input_size, hidden_size)
nn.LSTMCell(input_size, hidden_size)
nn.GRUCell(input_size, hidden_size)

# LSTM forward
# input: (seq_len, batch, input_size) if batch_first=False
#        (batch, seq_len, input_size) if batch_first=True
output, (h_n, c_n) = lstm(x)
# output: (batch, seq_len, hidden * directions) with batch_first=True
# h_n:   (num_layers * directions, batch, hidden)
# c_n:   (num_layers * directions, batch, hidden)

import torch.nn.functional as F

# q, k, v: (batch, num_heads, seq_len, head_dim)
out = F.scaled_dot_product_attention(q, k, v)                 # full attention
out = F.scaled_dot_product_attention(q, k, v, is_causal=True) # causal mask
out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask,
                                     dropout_p=0.1)
out = F.scaled_dot_product_attention(q, k, v, enable_gqa=True) # grouped-query

# Force a specific backend (e.g. to verify FlashAttention is used)
from torch.nn.attention import sdpa_kernel, SDPBackend

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
# Backends: FLASH_ATTENTION, EFFICIENT_ATTENTION, MATH, CUDNN_ATTENTION

# Multi-Head Attention
nn.MultiheadAttention(
    embed_dim=512,
    num_heads=8,
    dropout=0.0,
    bias=True,
    add_bias_kv=False,
    kdim=None,           # defaults to embed_dim
    vdim=None,
    batch_first=True,
)

# Usage
attn = nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)
out, weights = attn(query, key, value)
out, weights = attn(query, key, value, key_padding_mask=mask, attn_mask=causal_mask)

# Full Transformer
nn.Transformer(
    d_model=512, nhead=8,
    num_encoder_layers=6, num_decoder_layers=6,
    dim_feedforward=2048, dropout=0.1,
    activation='relu',   # or 'gelu' or callable
    batch_first=True,
)

# Individual encoder/decoder layers
nn.TransformerEncoderLayer(d_model=512, nhead=8, dim_feedforward=2048,
                           dropout=0.1, activation='gelu', batch_first=True,
                           norm_first=True)   # Pre-LN (more stable)
nn.TransformerEncoder(encoder_layer, num_layers=6, norm=nn.LayerNorm(512))

nn.TransformerDecoderLayer(d_model=512, nhead=8, batch_first=True)
nn.TransformerDecoder(decoder_layer, num_layers=6)

nn.Dropout(p=0.5)           # randomly zero elements
nn.Dropout2d(p=0.5)         # zero entire channels (for Conv2d feature maps)
nn.Dropout3d(p=0.5)
nn.AlphaDropout(p=0.5)      # SELU-compatible, preserves mean/variance

nn.ReLU(inplace=False)
nn.ReLU6()                  # min(max(0,x), 6)
nn.LeakyReLU(negative_slope=0.01)
nn.PReLU(num_parameters=1)  # learnable slope
nn.ELU(alpha=1.0)
nn.SELU()
nn.GELU()                   # Gaussian Error Linear Unit — default in transformers
nn.Mish()
nn.SiLU()                   # Swish: x * sigmoid(x)
nn.Sigmoid()
nn.Tanh()
nn.Softmax(dim=-1)
nn.LogSoftmax(dim=-1)
nn.Softplus(beta=1)
nn.Softsign()
nn.Hardswish()
nn.Hardsigmoid()
nn.Threshold(threshold=0.5, value=0.0)

nn.Embedding(num_embeddings=10000, embedding_dim=128, padding_idx=0,
             max_norm=None, scale_grad_by_freq=False, sparse=False)

# Usage
embed = nn.Embedding(10000, 128, padding_idx=0)
x = torch.randint(0, 10000, (8, 32))  # (batch, seq_len)
out = embed(x)                          # (8, 32, 128)

# Initialize from pre-trained (frozen)
embed = nn.Embedding.from_pretrained(pretrained_vectors, freeze=True, padding_idx=0)
# or copy into an existing layer without tracking gradients
with torch.no_grad():
    embed.weight.copy_(pretrained_vectors)
embed.weight.requires_grad = False       # freeze

# Embedding bag (aggregates embeddings without materializing all)
nn.EmbeddingBag(num_embeddings=10000, embedding_dim=128, mode='mean')

nn.ZeroPad2d(padding=(left, right, top, bottom))
nn.ConstantPad1d(padding=(1, 1), value=0.0)
nn.ConstantPad2d(padding=2, value=-1.0)
nn.ReflectionPad1d(padding=1)
nn.ReflectionPad2d(padding=(1, 1, 1, 1))
nn.ReplicationPad2d(padding=1)
nn.CircularPad2d(padding=2)

nn.Flatten(start_dim=1, end_dim=-1)        # flatten all dims except batch
nn.Unflatten(dim=1, unflattened_size=(2, 4))
nn.Identity()                               # pass-through (useful in ablations)
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
nn.Upsample(size=(224, 224), mode='nearest')
nn.PixelShuffle(upscale_factor=2)           # depth-to-space (ESPCN)
nn.PixelUnshuffle(downscale_factor=2)