Quantized layers module¶

This module provides drop-in quantized layer replacements for PyTorch models to support both floating-point and integer quantization-aware training (QAT).

All classes follow the same API as their original torch.nn counterparts. When a Pychop quantizer (with STE) is provided, weights and activations are fake-quantized during the forward pass while gradients flow through unchanged (Straight-Through Estimator).

Three STE-enabled quantizers are provided:

ChopSTE — floating-point (exponent + significand)
ChopfSTE — fixed-point (integer + fractional bits)
ChopiSTE — integer (uniform or symmetric)

STE quantizers (core)¶

class pychop.layers.ChopSTE(*args, **kwargs)[source]¶

Bases:

Create a ChopSTE instance for the current backend.

Raises¶

ImportError: If the current backend’s dependencies are not installed.

class pychop.layers.ChopfSTE(*args, **kwargs)[source]¶

Bases:

Create a ChopfSTE instance for the current backend.

Raises¶

ImportError: If the current backend’s dependencies are not installed.

class pychop.layers.ChopiSTE(*args, **kwargs)[source]¶

Bases:

Create a ChopiSTE instance for the current backend.

Raises¶

ImportError: If the current backend’s dependencies are not installed.

Utility Functions¶

pychop.layers.post_quantization(model, chop, eval_mode: bool = True, verbose: bool = False)[source]¶

Post-training quantization (PTQ) wrapper.

Dispatches to backend-specific implementation.

Parameters¶

modeltorch.nn.Module or flax.linen.Module: Neural network model.
chopChop, Chopf, or Chopi: Quantizer instance.
eval_modebool, default=True: Whether to set model to evaluation mode (PyTorch only).
verbosebool, default=False: Whether to print quantization details.

Returns¶

model: Quantized model.

Raises¶

ImportError: If the current backend’s dependencies are not installed.

Floating-point / Fixed-point quantized layers¶

These layers use ChopSTE (or Chop) for floating-point QAT.

class pychop.layers.QuantizedLinear(*args, **kwargs)¶

Bases:

Create a QuantizedLinear for the current backend.

class pychop.layers.QuantizedConv1d(*args, **kwargs)¶: Create a QuantizedConv1d for the current backend.

class pychop.layers.QuantizedConv2d(*args, **kwargs)¶: Create a QuantizedConv2d for the current backend.

class pychop.layers.QuantizedConv3d(*args, **kwargs)¶: Create a QuantizedConv3d for the current backend.

class pychop.layers.QuantizedConvTranspose1d(*args, **kwargs)¶: Create a QuantizedConvTranspose1d for the current backend.

class pychop.layers.QuantizedConvTranspose2d(*args, **kwargs)¶: Create a QuantizedConvTranspose2d for the current backend.

class pychop.layers.QuantizedConvTranspose3d(*args, **kwargs)¶: Create a QuantizedConvTranspose3d for the current backend.

class pychop.layers.QuantizedRNN(*args, **kwargs)¶: Create a QuantizedRNN for the current backend.

class pychop.layers.QuantizedLSTM(*args, **kwargs)¶: Create a QuantizedLSTM for the current backend.

class pychop.layers.QuantizedGRU(*args, **kwargs)¶: Create a QuantizedGRU for the current backend.

class pychop.layers.QuantizedMaxPool1d(*args, **kwargs)¶: Create a QuantizedMaxPool1d for the current backend.

class pychop.layers.QuantizedMaxPool2d(*args, **kwargs)¶: Create a QuantizedMaxPool2d for the current backend.

class pychop.layers.QuantizedMaxPool3d(*args, **kwargs)¶: Create a QuantizedMaxPool3d for the current backend.

class pychop.layers.QuantizedAvgPool1d(*args, **kwargs)¶: Create a QuantizedAvgPool1d for the current backend.

class pychop.layers.QuantizedAvgPool2d(*args, **kwargs)¶: Create a QuantizedAvgPool2d for the current backend.

class pychop.layers.QuantizedAvgPool3d(*args, **kwargs)¶: Create a QuantizedAvgPool3d for the current backend.

class pychop.layers.QuantizedAdaptiveAvgPool2d(*args, **kwargs)¶: Create a QuantizedAdaptiveAvgPool2d for the current backend.

class pychop.layers.QuantizedBatchNorm1d(*args, **kwargs)¶: Create a QuantizedBatchNorm1d for the current backend.

class pychop.layers.QuantizedBatchNorm2d(*args, **kwargs)¶: Create a QuantizedBatchNorm2d for the current backend.

class pychop.layers.QuantizedBatchNorm3d(*args, **kwargs)¶: Create a QuantizedBatchNorm3d for the current backend.

class pychop.layers.QuantizedLayerNorm(*args, **kwargs)¶: Create a QuantizedLayerNorm for the current backend.

class pychop.layers.QuantizedInstanceNorm1d(*args, **kwargs)¶: Create a QuantizedInstanceNorm1d for the current backend.

class pychop.layers.QuantizedInstanceNorm2d(*args, **kwargs)¶: Create a QuantizedInstanceNorm2d for the current backend.

class pychop.layers.QuantizedInstanceNorm3d(*args, **kwargs)¶: Create a QuantizedInstanceNorm3d for the current backend.

class pychop.layers.QuantizedGroupNorm(*args, **kwargs)¶: Create a QuantizedGroupNorm for the current backend.

class pychop.layers.QuantizedMultiheadAttention(*args, **kwargs)¶: Create a QuantizedMultiheadAttention for the current backend.

class pychop.layers.QuantizedEmbedding(*args, **kwargs)¶: Create a QuantizedEmbedding for the current backend.

Convenience aliases

pychop.layers.QuantizedAttention = <function _create_layer_factory.<locals>.factory>¶: Create a QuantizedMultiheadAttention for the current backend.

pychop.layers.QuantizedAvgPool = <function _create_layer_factory.<locals>.factory>¶: Create a QuantizedAvgPool2d for the current backend.

Activation & Regularization Layers (Floating-Point)¶

class pychop.layers.QuantizedReLU(*args, **kwargs)¶: Create a QuantizedReLU for the current backend.

class pychop.layers.QuantizedLeakyReLU(*args, **kwargs)¶: Create a QuantizedLeakyReLU for the current backend.

class pychop.layers.QuantizedSigmoid(*args, **kwargs)¶: Create a QuantizedSigmoid for the current backend.

class pychop.layers.QuantizedTanh(*args, **kwargs)¶: Create a QuantizedTanh for the current backend.

class pychop.layers.QuantizedGELU(*args, **kwargs)¶: Create a QuantizedGELU for the current backend.

class pychop.layers.QuantizedELU(*args, **kwargs)¶: Create a QuantizedELU for the current backend.

class pychop.layers.QuantizedPReLU(*args, **kwargs)¶: Create a QuantizedPReLU for the current backend.

class pychop.layers.QuantizedSoftmax(*args, **kwargs)¶: Create a QuantizedSoftmax for the current backend.

class pychop.layers.QuantizedDropout(*args, **kwargs)¶: Create a QuantizedDropout for the current backend.

Integer quantized layers¶

These layers use ChopiSTE for integer QAT (uniform or symmetric).

class pychop.layers.IQuantizedLinear(*args, **kwargs)¶: Create a IQuantizedLinear for the current backend.

class pychop.layers.IQuantizedConv1d(*args, **kwargs)¶: Create a IQuantizedConv1d for the current backend.

class pychop.layers.IQuantizedConv2d(*args, **kwargs)¶: Create a IQuantizedConv2d for the current backend.

class pychop.layers.IQuantizedConv3d(*args, **kwargs)¶: Create a IQuantizedConv3d for the current backend.

class pychop.layers.IQuantizedConvTranspose1d(*args, **kwargs)¶: Create a IQuantizedConvTranspose1d for the current backend.

class pychop.layers.IQuantizedConvTranspose2d(*args, **kwargs)¶: Create a IQuantizedConvTranspose2d for the current backend.

class pychop.layers.IQuantizedConvTranspose3d(*args, **kwargs)¶: Create a IQuantizedConvTranspose3d for the current backend.

class pychop.layers.IQuantizedRNN(*args, **kwargs)¶: Create a IQuantizedRNN for the current backend.

class pychop.layers.IQuantizedLSTM(*args, **kwargs)¶: Create a IQuantizedLSTM for the current backend.

class pychop.layers.IQuantizedGRU(*args, **kwargs)¶: Create a IQuantizedGRU for the current backend.

class pychop.layers.IQuantizedMaxPool1d(*args, **kwargs)¶: Create a IQuantizedMaxPool1d for the current backend.

class pychop.layers.IQuantizedMaxPool2d(*args, **kwargs)¶: Create a IQuantizedMaxPool2d for the current backend.

class pychop.layers.IQuantizedMaxPool3d(*args, **kwargs)¶: Create a IQuantizedMaxPool3d for the current backend.

class pychop.layers.IQuantizedAvgPool1d(*args, **kwargs)¶: Create a IQuantizedAvgPool1d for the current backend.

class pychop.layers.IQuantizedAvgPool2d(*args, **kwargs)¶: Create a IQuantizedAvgPool2d for the current backend.

class pychop.layers.IQuantizedAvgPool3d(*args, **kwargs)¶: Create a IQuantizedAvgPool3d for the current backend.

class pychop.layers.IQuantizedAdaptiveAvgPool1d(*args, **kwargs)¶: Create a IQuantizedAdaptiveAvgPool1d for the current backend.

class pychop.layers.IQuantizedAdaptiveAvgPool2d(*args, **kwargs)¶: Create a IQuantizedAdaptiveAvgPool2d for the current backend.

class pychop.layers.IQuantizedAdaptiveAvgPool3d(*args, **kwargs)¶: Create a IQuantizedAdaptiveAvgPool3d for the current backend.

class pychop.layers.IQuantizedBatchNorm1d(*args, **kwargs)¶: Create a IQuantizedBatchNorm1d for the current backend.

class pychop.layers.IQuantizedBatchNorm2d(*args, **kwargs)¶: Create a IQuantizedBatchNorm2d for the current backend.

class pychop.layers.IQuantizedBatchNorm3d(*args, **kwargs)¶: Create a IQuantizedBatchNorm3d for the current backend.

class pychop.layers.IQuantizedLayerNorm(*args, **kwargs)¶: Create a IQuantizedLayerNorm for the current backend.

class pychop.layers.IQuantizedInstanceNorm1d(*args, **kwargs)¶: Create a IQuantizedInstanceNorm1d for the current backend.

class pychop.layers.IQuantizedInstanceNorm2d(*args, **kwargs)¶: Create a IQuantizedInstanceNorm2d for the current backend.

class pychop.layers.IQuantizedInstanceNorm3d(*args, **kwargs)¶: Create a IQuantizedInstanceNorm3d for the current backend.

class pychop.layers.IQuantizedGroupNorm(*args, **kwargs)¶: Create a IQuantizedGroupNorm for the current backend.

class pychop.layers.IQuantizedMultiheadAttention(*args, **kwargs)¶: Create a IQuantizedMultiheadAttention for the current backend.

class pychop.layers.IQuantizedEmbedding(*args, **kwargs)¶: Create a IQuantizedEmbedding for the current backend.

Integer activation & regularization layers¶

class pychop.layers.IQuantizedReLU(*args, **kwargs)¶: Create a IQuantizedReLU for the current backend.

class pychop.layers.IQuantizedLeakyReLU(*args, **kwargs)¶: Create a IQuantizedLeakyReLU for the current backend.

class pychop.layers.IQuantizedSigmoid(*args, **kwargs)¶: Create a IQuantizedSigmoid for the current backend.

class pychop.layers.IQuantizedTanh(*args, **kwargs)¶: Create a IQuantizedTanh for the current backend.

class pychop.layers.IQuantizedGELU(*args, **kwargs)¶: Create a IQuantizedGELU for the current backend.

class pychop.layers.IQuantizedELU(*args, **kwargs)¶: Create a IQuantizedELU for the current backend.

class pychop.layers.IQuantizedSiLU(*args, **kwargs)¶: Create a IQuantizedSiLU for the current backend.

class pychop.layers.IQuantizedPReLU(*args, **kwargs)¶: Create a IQuantizedPReLU for the current backend.

class pychop.layers.IQuantizedSoftmax(*args, **kwargs)¶: Create a IQuantizedSoftmax for the current backend.

class pychop.layers.IQuantizedDropout(*args, **kwargs)¶: Create a IQuantizedDropout for the current backend.

Convenience aliases

pychop.layers.IQuantizedAttention = <function _create_layer_factory.<locals>.factory>¶: Create a IQuantizedMultiheadAttention for the current backend.

pychop.layers.IQuantizedAvgPool = <function _create_layer_factory.<locals>.factory>¶: Create a IQuantizedAvgPool2d for the current backend.

Usage Example¶

from pychop.layers import (
    ChopSTE, ChopfSTE, ChopiSTE,
    QuantizedConv2d, QuantizedReLU,
    IQuantizedLinear, IQuantizedReLU
)

# Floating-point QAT
chop_fp = ChopSTE(exp_bits=5, sig_bits=10, rmode=3)

# Fixed-point QAT
chop_fixed = ChopfSTE(ibits=8, fbits=8, rmode=1)

# Integer QAT
chop_int = ChopiSTE(bits=8, symmetric=True)

class MyQuantizedNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = QuantizedConv2d(3, 64, 3, chop=chop_fp)
        self.relu  = QuantizedReLU(chop=chop_fp)
        self.fc    = IQuantizedLinear(512, 10, chop=chop_int)

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = x.view(x.size(0), -1)
        return self.fc(x)

Post-training quantization (PTQ)¶

from pychop.layers import post_quantization, ChopSTE

chop = ChopSTE(exp_bits=8, sig_bits=23)   # or any other chop
quantized_model = post_quantization(model, chop, eval_mode=True, verbose=True)