Quantized layers module

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

For PyTorch, all classes follow the same API as their original torch.nn counterparts. For TensorFlow, all classes follow the same API as their original tf.keras.layers 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)

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)

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)

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)

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)

TensorFlow / Keras Quantized Layers

Pychop also provides quantized layer replacements for TensorFlow Keras models via the pychop.tf.layers module. These layers follow the same API as their original tf.keras.layers counterparts and support STE-based quantization-aware training.

Available layers include:

• Linear: QuantizedLinear (Dense replacement)

• Convolution: QuantizedConv1d, QuantizedConv2d, QuantizedConv3d

• Pooling: QuantizedMaxPool1d, QuantizedMaxPool2d, QuantizedAvgPool1d, QuantizedAvgPool2d, QuantizedAdaptiveAvgPool2d

• Normalization: QuantizedBatchNorm1d, QuantizedBatchNorm2d, QuantizedLayerNorm, QuantizedInstanceNorm1d, QuantizedInstanceNorm2d, QuantizedGroupNorm

• Activation: QuantizedReLU, QuantizedLeakyReLU, QuantizedSigmoid, QuantizedTanh, QuantizedGELU, QuantizedELU, QuantizedSiLU, QuantizedPReLU

• Other: QuantizedEmbedding, QuantizedMultiheadAttention, QuantizedDropout, QuantizedSoftmax

Usage Example (TensorFlow):

import tensorflow as tf
import pychop
from pychop.tf.layers import (
    ChopSTE, QuantizedConv2d, QuantizedLinear, QuantizedReLU
)

pychop.backend('tensorflow')

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

class MyQuantizedNet(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.conv1 = QuantizedConv2d(64, 3, chop=chop_fp)
        self.relu  = QuantizedReLU(chop=chop_fp)
        self.fc    = QuantizedLinear(10, chop=chop_fp)

    def call(self, x):
        x = self.relu(self.conv1(x))
        x = tf.reshape(x, [tf.shape(x)[0], -1])
        return self.fc(x)