Supported PyTorch Operations

Native PyTorch operations that can be traced, quantized, and lowered.

Layers

PyTorch Op	Notes
nn.Conv2d	Dense convolution (groups=1); depthwise only when groups = in_channels = out_channels (depth multiplier 1 only)
nn.Linear
nn.AvgPool2d	A preceding F.pad folds into its padding
nn.MaxPool2d	No padding support, native or via a preceding F.pad

Activations

PyTorch Op	Notes
torch.softmax / F.softmax	Output pinned to a fixed quantization grid
torch.sigmoid	Output pinned to a fixed quantization grid
F.relu	Only recognized as a fusion tail immediately after a weight op (Conv2d/Linear) or a residual add; not a standalone layer

Shape & Elementwise Operations

PyTorch Op	Notes
torch.flatten
torch.reshape / view	Cannot change the batch-axis extent
torch.transpose	Feature rank 2–3 only; cannot move the batch axis
torch.slice / __getitem__[:n]	Sequence-domain tensors only (not spatial/HWC feature maps); not the batch axis; step must be 1
F.pad(pad_value=0)	Only usable as an immediate prefix to Conv2d or AvgPool2d; no standalone padding layer in quantized models
torch.cat	Not the batch axis; requires compiler ≥ v0.15.0; inputs and output share one quantization scale
torch.matmul / torch.bmm / torch.mm	Activation × activation; feature rank 2–3 per operand; trailing scale must be a scalar constant; output must feed a quantized op
torch.mul(x, const_tensor)	Activation × constant-tensor broadcast only (e.g., RoPE's cos/sin tables); not a bare scalar, not activation × activation
Element-wise +	Both operands must already be quantized activation streams
F.interpolate(mode="nearest")	Mode must be "nearest" or "nearest-exact"; input/output share one quantization grid

Not Supported

These operations look like they should work but are not supported. Using any of them raises an error during export, so they will not pass silently into a broken model.

Operation	Why
torch.permute	Only the NCHW→NHWC reorder inside dg.Flatten is recognized; any other permute is rejected. Use dg.Flatten or torch.transpose.
Standalone F.pad	Padding is only supported immediately before a Conv2d or AvgPool2d.
Tensor.expand()	Not supported.
Tensor.repeat()	Only supported inside the logic layers (dg.BitShift, etc.), not in quantized models.
Bare torch.mul(x, scalar)	Not supported on its own; only valid inside dg.Norm or an attention score scale.
Standalone F.relu	Only supported as a fusion tail after a weight op or residual add, not as a layer of its own.
nn.ConvTranspose2d	Not supported.
Other activations	F.gelu, torch.tanh, F.silu, F.hardswish, F.leaky_relu, and similar are not supported; only relu, sigmoid, and softmax are.
Fused attention	F.scaled_dot_product_attention and nn.MultiheadAttention export as a single op that cannot be quantized. Use dg.SelfAttention or explicit matmul + softmax.
Dilated nn.Conv2d	Dilation (dilation != 1) is not supported, and is not caught at export, so avoid it.

Special Cases

Operation	Behavior
Batch normalization	Conv → BatchNorm folds via torchao's QAT pass inside freeze_quantization; Linear → BatchNorm folds separately inside dg.export() (requires the linear have exactly one user)
Dropout	Absent from the traced graph in eval mode