chore: import upstream snapshot with attribution

tools/pytorch-quantization/pytorch_quantization/nn/modules/_utils.py

@@ -0,0 +1,165 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+"""Some helper functions for implementing quantized modules"""
+import copy
+import inspect
+
+from absl import logging
+
+from torch import nn
+
+from pytorch_quantization.nn import TensorQuantizer
+from pytorch_quantization.tensor_quant import QuantDescriptor, QUANT_DESC_8BIT_PER_TENSOR
+
+
+class QuantMixin():
+    """Mixin class for adding basic quantization logic to quantized modules"""
+
+    default_quant_desc_input = QUANT_DESC_8BIT_PER_TENSOR
+    default_quant_desc_weight = QUANT_DESC_8BIT_PER_TENSOR
+
+    @classmethod
+    def set_default_quant_desc_input(cls, value):
+        """
+        Args:
+            value: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+        """
+        if not isinstance(value, QuantDescriptor):
+            raise ValueError("{} is not an instance of QuantDescriptor!")
+        cls.default_quant_desc_input = copy.deepcopy(value)
+
+    @classmethod
+    def set_default_quant_desc_weight(cls, value):
+        """
+        Args:
+            value: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+        """
+        if not isinstance(value, QuantDescriptor):
+            raise ValueError("{} is not an instance of QuantDescriptor!")
+        cls.default_quant_desc_weight = copy.deepcopy(value)
+
+    def init_quantizer(self, quant_desc_input, quant_desc_weight, num_layers=None):
+        """Helper function for __init__ of quantized module
+
+        Create input and weight quantizer based on quant_desc passed by kwargs, or default of the class.
+
+        Args:
+            quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+            quant_desc_weight: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+            num_layers: An integer. Default None. If not None, create a list of quantizers.
+        """
+        if not inspect.stack()[1].function == "__init__":
+            raise TypeError("{} should be only called by __init__ of quantized module.".format(__name__))
+        self._fake_quant = True
+        if (not quant_desc_input.fake_quant) or (not quant_desc_weight.fake_quant):
+            raise ValueError("Only fake quantization is supported!")
+
+        logging.info("Input is %squantized to %d bits in %s with axis %s!", ""
+                     if not quant_desc_input.fake_quant else "fake ",
+                     quant_desc_input.num_bits, self.__class__.__name__, quant_desc_input.axis)
+        logging.info("Weight is %squantized to %d bits in %s with axis %s!", ""
+                     if not quant_desc_weight.fake_quant else "fake ",
+                     quant_desc_weight.num_bits, self.__class__.__name__, quant_desc_weight.axis)
+
+        if num_layers is None:
+            self._input_quantizer = TensorQuantizer(quant_desc_input)
+            self._weight_quantizer = TensorQuantizer(quant_desc_weight)
+        else:
+            self._input_quantizers = nn.ModuleList([TensorQuantizer(quant_desc_input) for _ in range(num_layers)])
+            self._weight_quantizers = nn.ModuleList([TensorQuantizer(quant_desc_weight) for _ in range(num_layers)])
+
+    # pylint:disable=missing-docstring
+    @property
+    def input_quantizer(self):
+        return self._input_quantizer
+
+    @property
+    def weight_quantizer(self):
+        return self._weight_quantizer
+    # pylint:enable=missing-docstring
+
+
+class QuantInputMixin():
+    """Mixin class for adding basic quantization logic to quantized modules"""
+
+    default_quant_desc_input = QUANT_DESC_8BIT_PER_TENSOR
+
+    @classmethod
+    def set_default_quant_desc_input(cls, value):
+        """
+        Args:
+            value: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+        """
+        if not isinstance(value, QuantDescriptor):
+            raise ValueError("{} is not an instance of QuantDescriptor!")
+        cls.default_quant_desc_input = copy.deepcopy(value)
+
+    def init_quantizer(self, quant_desc_input):
+        """Helper function for __init__ of simple quantized module
+
+        Create input quantizer based on quant_desc passed by kwargs, or default of the class.
+
+        Args:
+            quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+        """
+        if not inspect.stack()[1].function == "__init__":
+            raise TypeError("{} should be only called by __init__ of quantized module.".format(__name__))
+        self._fake_quant = True
+        if not quant_desc_input.fake_quant:
+            raise ValueError("Only fake quantization is supported!")
+
+        logging.info("Input is %squantized to %d bits in %s with axis %s!", ""
+                     if not quant_desc_input.fake_quant else "fake ",
+                     quant_desc_input.num_bits, self.__class__.__name__, quant_desc_input.axis)
+
+        self._input_quantizer = TensorQuantizer(quant_desc_input)
+
+    # pylint:disable=missing-docstring
+    @property
+    def input_quantizer(self):
+        return self._input_quantizer
+    # pylint:enable=missing-docstring
+
+
+def pop_quant_desc_in_kwargs(quant_cls, input_only=False, **kwargs):
+    """Pop quant descriptors in kwargs
+
+    If there is no descriptor in kwargs, the default one in quant_cls will be used
+
+    Arguments:
+       quant_cls: A class that has default quantization descriptors
+       input_only: A boolean. If True, pop quant_desc_input only, not quant_desc_weight. Default false.
+
+    Keyword Arguments:
+       quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+           Quantization descriptor of input.
+       quant_desc_weight: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+           Quantization descriptor of weight.
+    """
+    quant_desc_input = kwargs.pop('quant_desc_input', quant_cls.default_quant_desc_input)
+    if not input_only:
+        quant_desc_weight = kwargs.pop('quant_desc_weight', quant_cls.default_quant_desc_weight)
+
+    # Check if anything is left in **kwargs
+    if kwargs:
+        raise TypeError("Unused keys: {}".format(kwargs.keys()))
+
+    if input_only:
+        return quant_desc_input
+    return quant_desc_input, quant_desc_weight

tools/pytorch-quantization/pytorch_quantization/nn/modules/clip.py

@@ -0,0 +1,59 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+"""Implement a clip module as pytorch only has a simple clamp function """
+import torch
+from torch import nn
+from torch.nn.parameter import Parameter
+
+from pytorch_quantization.nn import functional as QF
+
+__all__ = ['Clip']
+
+class Clip(nn.Module):
+    """Clip tensor
+
+    Args:
+        clip_value_min: A number or tensor of lower bound to clip
+        clip_value_max: A number of tensor of upper bound to clip
+        learn_min: A boolean. If True, learn min. clip_value_min will be used to initialize. Default False
+        learn_max: A boolean. Similar as learn_min but for max.
+
+    Raises:
+        ValueError:
+    """
+
+    def __init__(self, clip_value_min, clip_value_max, learn_min=False, learn_max=False):
+        super(Clip, self).__init__()
+        if learn_min:
+            if not isinstance(clip_value_min, float) and clip_value_min.size != 1:
+                raise ValueError("clip_value_min/clip_value_max must be scalar for initilizing learnable range.")
+            self.clip_value_min = Parameter(torch.tensor(clip_value_min))  # pylint: disable=not-callable
+        else:
+            self.clip_value_min = clip_value_min
+
+        if learn_max:
+            if not isinstance(clip_value_max, float) and clip_value_max.size != 1:
+                raise ValueError("clip_value_min/clip_value_max must be scalar for initilizing learnable range.")
+            self.clip_value_max = Parameter(torch.tensor(clip_value_max))  # pylint: disable=not-callable
+        else:
+            self.clip_value_max = clip_value_max
+
+    def forward(self, inputs):
+        outputs = QF.clip(inputs, self.clip_value_min, self.clip_value_max)
+        return outputs

tools/pytorch-quantization/pytorch_quantization/nn/modules/quant_conv.py

@@ -0,0 +1,419 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+"""Quantized convolution
+Base code is from nn.Conv, details of Module and original argument can be found there.
+Module names are intentionally kept same as unquantized version so that they can be dropped into preexisting model
+easily, and load pretrained weight. Aliases with Quant prefix are defined and are encouraged to be used explicitly
+when start scratch.
+"""
+
+import inspect
+import torch
+import torch.nn
+import torch.nn.functional as F
+from torch.nn.modules.utils import _single, _pair, _triple
+from torch.nn.modules.conv import _ConvTransposeNd
+
+from pytorch_quantization import tensor_quant
+
+from . import _utils
+
+__all__ = [
+    "Conv2d", "QuantConv2d", "Conv3d", "QuantConv3d", "Conv1d", "QuantConv1d", "ConvTranspose1d", "ConvTranspose2d",
+    "ConvTranspose3d", "QuantConvTranspose1d", "QuantConvTranspose2d", "QuantConvTranspose3d"
+]
+
+
+class _QuantConvNd(torch.nn.modules.conv._ConvNd, _utils.QuantMixin):
+    """base class of quantized Conv inherited from _ConvNd
+
+    Comments of original arguments can be found in torch.nn.modules.conv
+
+    Arguments:
+        quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of input.
+        quant_desc_weight: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of weight.
+
+    Raises:
+        ValueError: If unsupported arguments are passed in.
+
+    Readonly properties:
+        - input_quantizer:
+        - weight_quantizer:
+
+    Static methods:
+        - set_default_quant_desc_input: Set default_quant_desc_input
+        - set_default_quant_desc_weight: Set default_quant_desc_weight
+    """
+
+    default_quant_desc_input = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, transposed, output_padding,
+                 groups, bias, padding_mode, quant_desc_input, quant_desc_weight):
+        super(_QuantConvNd, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation,
+                                           transposed, output_padding, groups, bias, padding_mode)
+        self.init_quantizer(quant_desc_input, quant_desc_weight)
+
+    def _quant(self, input):
+        """Apply quantization on input and weight
+
+        Function called by the classes lower in the hierarchy, which actually performs the quantization before forward
+        in the derivate class the particular Function.
+
+        Arguments:
+            input: in_features to quantize
+        Returns:
+            A tuple: (quant_in_feature, quant_weight)
+        """
+        quant_input = self._input_quantizer(input)
+        quant_weight = self._weight_quantizer(self.weight)
+
+        return (quant_input, quant_weight)
+
+
+class QuantConv2d(_QuantConvNd):
+    """Quantized 2D conv"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias=True,
+                 padding_mode='zeros',
+                 **kwargs):
+
+        kernel_size = _pair(kernel_size)
+        stride = _pair(stride)
+        padding = _pair(padding)
+        dilation = _pair(dilation)
+
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConv2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, False,
+                                          _pair(0), groups, bias, padding_mode,
+                                          quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input):
+        # the actual quantization happens in the next level of the class hierarchy
+        quant_input, quant_weight = self._quant(input)
+
+        if self.padding_mode == 'circular':
+            expanded_padding = ((self.padding[1] + 1) // 2, self.padding[1] // 2,
+                                (self.padding[0] + 1) // 2, self.padding[0] // 2)
+            output = F.conv2d(F.pad(quant_input, expanded_padding, mode='circular'),
+                              quant_weight, self.bias, self.stride,
+                              _pair(0), self.dilation, self.groups)
+        else:
+            output = F.conv2d(quant_input, quant_weight, self.bias, self.stride, self.padding, self.dilation,
+                              self.groups)
+
+        return output
+
+
+class QuantConv3d(_QuantConvNd):
+    """Quantized 3D Conv"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias=True,
+                 padding_mode='zeros',
+                 **kwargs):
+
+        kernel_size = _triple(kernel_size)
+        stride = _triple(stride)
+        padding = _triple(padding)
+        dilation = _triple(dilation)
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConv3d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, False,
+                                          _triple(0), groups, bias, padding_mode,
+                                          quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input):
+        # the actual quantization happens in the next level of the class hierarchy
+        quant_input, quant_weight = self._quant(input)
+
+        if self.padding_mode == 'circular':
+            expanded_padding = ((self.padding[2] + 1) // 2, self.padding[2] // 2,
+                                (self.padding[1] + 1) // 2, self.padding[1] // 2,
+                                (self.padding[0] + 1) // 2, self.padding[0] // 2)
+            output = F.conv3d(F.pad(quant_input, expanded_padding, mode='circular'),
+                              quant_weight, self.bias, self.stride, _triple(0),
+                              self.dilation, self.groups)
+        else:
+            output = F.conv3d(quant_input, quant_weight, self.bias, self.stride, self.padding, self.dilation,
+                              self.groups)
+
+        return output
+
+
+class QuantConv1d(_QuantConvNd):
+    """Quantized 1D Conv"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV1D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias=True,
+                 padding_mode='zeros',
+                 **kwargs):
+
+        kernel_size = _single(kernel_size)
+        stride = _single(stride)
+        padding = _single(padding)
+        dilation = _single(dilation)
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConv1d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, False,
+                                          _single(0), groups, bias, padding_mode,
+                                          quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input):
+        # the actual quantization happens in the next level of the class hierarchy
+        quant_input, quant_weight = self._quant(input)
+
+        if self.padding_mode == 'circular':
+            expanded_padding = ((self.padding[0] + 1) // 2, self.padding[0] // 2)
+            output = F.conv1d(F.pad(quant_input, expanded_padding, mode='circular'),
+                              quant_weight, self.bias, self.stride,
+                              _single(0), self.dilation, self.groups)
+        else:
+            output = F.conv1d(quant_input, quant_weight, self.bias, self.stride,
+                              self.padding, self.dilation, self.groups)
+
+        return output
+
+
+class _QuantConvTransposeNd(torch.nn.modules.conv._ConvTransposeNd, _utils.QuantMixin):
+    """base class of quantized Transposed Conv inherited from _ConvTransposeNd
+
+    Comments of original arguments can be found in torch.nn.modules.conv
+
+    Arguments:
+        quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of input.
+        quant_desc_weight: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of weight.
+
+    Raises:
+        ValueError: If unsupported arguments are passed in.
+
+    Readonly properties:
+        - input_quantizer:
+        - weight_quantizer:
+
+    Static methods:
+        - set_default_quant_desc_input: Set default_quant_desc_input
+        - set_default_quant_desc_weight: Set default_quant_desc_weight
+    """
+
+    default_quant_desc_input = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+    def __init__(self, in_channels, out_channels, kernel_size, stride,
+                 padding, dilation, transposed, output_padding,
+                 groups, bias, padding_mode, quant_desc_input, quant_desc_weight):
+        super(_QuantConvTransposeNd, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation,
+                                                    transposed, output_padding, groups, bias, padding_mode)
+        self.init_quantizer(quant_desc_input, quant_desc_weight)
+
+    def _quant(self, input):
+        """Apply quantization on input and weight
+
+        Function called by the classes lower in the hierarchy, which actually performs the quantization before forward
+        in the derivate class the particular Function.
+
+        Arguments:
+            input: in_features to quantize
+        Returns:
+            A tuple: (quant_in_feature, quant_weight)
+        """
+        quant_input = self._input_quantizer(input)
+        quant_weight = self._weight_quantizer(self.weight)
+
+        return (quant_input, quant_weight)
+
+    def _output_padding_nd(self,
+                           input,
+                           output_size,
+                           stride,
+                           padding,
+                           kernel_size,
+                           num_spatial_dims,
+                           dilation=None):
+        if "num_spatial_dims" in inspect.signature(self._output_padding).parameters:
+            return self._output_padding(input, output_size, stride, padding, kernel_size, num_spatial_dims)
+        else:
+            return self._output_padding(input, output_size, stride, padding, kernel_size)
+
+
+class QuantConvTranspose1d(_QuantConvTransposeNd):
+    """Quantized ConvTranspose1d"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 output_padding=0,
+                 groups=1,
+                 bias=True,
+                 dilation=1,
+                 padding_mode='zeros',
+                 **kwargs):
+        kernel_size = _single(kernel_size)
+        stride = _single(stride)
+        padding = _single(padding)
+        dilation = _single(dilation)
+        output_padding = _single(output_padding)
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConvTranspose1d, self).__init__(
+            in_channels, out_channels, kernel_size, stride, padding, dilation,
+            True, output_padding, groups, bias, padding_mode,
+            quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input, output_size=None):
+        if self.padding_mode != 'zeros':
+            raise ValueError('Only `zeros` padding mode is supported for QuantConvTranspose1d')
+
+        num_spatial_dims = 1
+        output_padding = self._output_padding_nd(input, output_size, self.stride, self.padding, self.kernel_size,
+                                                 num_spatial_dims)
+
+        quant_input, quant_weight = self._quant(input)
+        output = F.conv_transpose1d(quant_input, quant_weight, self.bias, self.stride, self.padding, output_padding,
+                                    self.groups, self.dilation)
+        return output
+
+
+class QuantConvTranspose2d(_QuantConvTransposeNd):
+    """Quantized ConvTranspose2d"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 output_padding=0,
+                 groups=1,
+                 bias=True,
+                 dilation=1,
+                 padding_mode='zeros',
+                 **kwargs):
+        kernel_size = _pair(kernel_size)
+        stride = _pair(stride)
+        padding = _pair(padding)
+        dilation = _pair(dilation)
+        output_padding = _pair(output_padding)
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConvTranspose2d, self).__init__(
+            in_channels, out_channels, kernel_size, stride, padding, dilation,
+            True, output_padding, groups, bias, padding_mode,
+            quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input, output_size=None):
+        if self.padding_mode != 'zeros':
+            raise ValueError('Only `zeros` padding mode is supported for QuantConvTranspose2d')
+
+        num_spatial_dims = 2
+        output_padding = self._output_padding_nd(input, output_size, self.stride, self.padding, self.kernel_size,
+                                                 num_spatial_dims)
+
+        quant_input, quant_weight = self._quant(input)
+        output = F.conv_transpose2d(quant_input, quant_weight, self.bias, self.stride, self.padding, output_padding,
+                                    self.groups, self.dilation)
+
+        return output
+
+
+class QuantConvTranspose3d(_QuantConvTransposeNd):
+    """Quantized ConvTranspose3d"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 output_padding=0,
+                 groups=1,
+                 bias=True,
+                 dilation=1,
+                 padding_mode='zeros',
+                 **kwargs):
+        kernel_size = _triple(kernel_size)
+        stride = _triple(stride)
+        padding = _triple(padding)
+        dilation = _triple(dilation)
+        output_padding = _triple(output_padding)
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConvTranspose3d, self).__init__(
+            in_channels, out_channels, kernel_size, stride, padding, dilation,
+            True, output_padding, groups, bias, padding_mode,
+            quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input, output_size=None):
+        if self.padding_mode != 'zeros':
+            raise ValueError('Only `zeros` padding mode is supported for QuantConvTranspose3d')
+
+        num_spatial_dims = 3
+        output_padding = self._output_padding_nd(input, output_size, self.stride, self.padding, self.kernel_size,
+                                                 num_spatial_dims)
+
+        quant_input, quant_weight = self._quant(input)
+        output = F.conv_transpose3d(quant_input, quant_weight, self.bias, self.stride, self.padding, output_padding,
+                                    self.groups, self.dilation)
+
+        return output
+
+
+# Define alias with Quant prefix
+_ConvNd = _QuantConvNd
+Conv1d = QuantConv1d
+Conv2d = QuantConv2d
+Conv3d = QuantConv3d
+ConvTranspose1d = QuantConvTranspose1d
+ConvTranspose2d = QuantConvTranspose2d
+ConvTranspose3d = QuantConvTranspose3d

tools/pytorch-quantization/pytorch_quantization/nn/modules/quant_instancenorm.py

@@ -0,0 +1,79 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Quantized instance normalization module
+   Base code is from nn.InstanceNorm, details of the module can be found from the offical repo.
+"""
+
+from torch.nn.modules.batchnorm import _NormBase
+import torch.nn.functional as F
+from torch.nn.modules import instancenorm
+
+from pytorch_quantization.nn import TensorQuantizer
+from pytorch_quantization import tensor_quant
+from . import _utils
+
+__all__ = [
+    "QuantInstanceNorm1d", "QuantInstanceNorm2d", "QuantInstanceNorm3d"
+]
+
+class QuantInstanceNorm1d(instancenorm.InstanceNorm1d, _utils.QuantInputMixin):
+    r"""Applies Quantized Instance Normalization over a 3D input
+    """
+    def __init__(
+            self, num_features: int, eps: float = 1e-5, momentum: float = 0.1, affine: bool = False,
+            track_running_stats: bool = False, **kwargs):
+        super(QuantInstanceNorm1d, self).__init__(
+            num_features, eps, momentum, affine, track_running_stats)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantInstanceNorm1d, self).forward(quant_input)
+
+
+class QuantInstanceNorm2d(instancenorm.InstanceNorm2d, _utils.QuantInputMixin):
+    r"""Applies Quantized Instance Normalization over a 4D input
+    """
+    def __init__(
+            self, num_features: int, eps: float = 1e-5, momentum: float = 0.1, affine: bool = False,
+            track_running_stats: bool = False, **kwargs):
+        super(QuantInstanceNorm2d, self).__init__(
+            num_features, eps, momentum, affine, track_running_stats)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantInstanceNorm2d, self).forward(quant_input)
+
+
+class QuantInstanceNorm3d(instancenorm.InstanceNorm3d, _utils.QuantInputMixin):
+    r"""Applies Quantized Instance Normalization over a 5D input
+    """
+    def __init__(
+            self, num_features: int, eps: float = 1e-5, momentum: float = 0.1, affine: bool = False,
+            track_running_stats: bool = False, **kwargs):
+        super(QuantInstanceNorm3d, self).__init__(
+            num_features, eps, momentum, affine, track_running_stats)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantInstanceNorm3d, self).forward(quant_input)

tools/pytorch-quantization/pytorch_quantization/nn/modules/quant_linear.py

@@ -0,0 +1,78 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+"""Quantized Linear"""
+from torch import nn
+from torch.nn import functional as F
+
+from pytorch_quantization import tensor_quant
+
+from . import _utils
+
+__all__ = ["Linear", "QuantLinear"]
+
+class QuantLinear(nn.Linear, _utils.QuantMixin):
+    """Quantized version of nn.Linear
+
+    Apply quantized linear to the incoming data, y = dequant(quant(x)quant(A)^T + b).
+
+    Keep Module name "Linear" instead of "QuantLinear" so that it can be easily dropped into preexisting model and load
+    pretrained weights. An alias "QuantLinear" is defined below. The base code is a copy of nn.Linear, see detailed
+    comment of original arguments there.
+
+    Quantization descriptors are passed in in kwargs. If not presents, default_quant_desc_input and
+    default_quant_desc_weight are used.
+
+    Keyword Arguments:
+        quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of input.
+        quant_desc_wegiht: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of weight.
+
+    Raises:
+        ValueError: If unsupported arguments are passed in.
+        KeyError: If unsupported kwargs are passed in.
+
+    Readonly properties:
+        - input_quantizer:
+        - weight_quantizer:
+
+    Static methods:
+        - set_default_quant_desc_input: Set default_quant_desc_input
+        - set_default_quant_desc_weight: Set default_quant_desc_weight
+    """
+
+    default_quant_desc_input = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW
+
+    def __init__(self, in_features, out_features, bias=True, **kwargs):
+        super(QuantLinear, self).__init__(in_features, out_features, bias)
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+
+        self.init_quantizer(quant_desc_input, quant_desc_weight)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        quant_weight = self._weight_quantizer(self.weight)
+
+        output = F.linear(quant_input, quant_weight, bias=self.bias)
+
+        return output
+
+
+Linear = QuantLinear

tools/pytorch-quantization/pytorch_quantization/nn/modules/quant_pooling.py

@@ -0,0 +1,163 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+"""Quantized Pooling
+Base code is from nn.pooling, details of Module and original argument can be found there.
+Module names are intentionally kept same as unquantized version so that they can be dropped into preexisting model
+easily, and load pretrained weight. Aliases with Quant prefix are defined and are encouraged to be used explicitly
+when start scratch.
+"""
+
+from torch.nn.modules import pooling
+
+from . import _utils
+
+__all__ = [
+    "MaxPool1d", "QuantMaxPool1d", "MaxPool2d", "QuantMaxPool2d", "MaxPool3d", "QuantMaxPool3d",
+    "AvgPool1d", "QuantAvgPool1d", "AvgPool2d", "QuantAvgPool2d", "AvgPool3d", "QuantAvgPool3d",
+    "AdaptiveAvgPool1d", "QuantAdaptiveAvgPool1d", "AdaptiveAvgPool2d", "QuantAdaptiveAvgPool2d",
+    "AdaptiveAvgPool3d", "QuantAdaptiveAvgPool3d"
+]
+
+class QuantMaxPool1d(pooling.MaxPool1d, _utils.QuantInputMixin):
+    """Quantized 1D maxpool"""
+    def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
+                 return_indices=False, ceil_mode=False, **kwargs):
+        super(QuantMaxPool1d, self).__init__(kernel_size, stride, padding, dilation,
+                                             return_indices, ceil_mode)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantMaxPool1d, self).forward(quant_input)
+
+class QuantMaxPool2d(pooling.MaxPool2d, _utils.QuantInputMixin):
+    """Quantized 2D maxpool"""
+    def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
+                 return_indices=False, ceil_mode=False, **kwargs):
+        super(QuantMaxPool2d, self).__init__(kernel_size, stride, padding, dilation,
+                                             return_indices, ceil_mode)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantMaxPool2d, self).forward(quant_input)
+
+class QuantMaxPool3d(pooling.MaxPool3d, _utils.QuantInputMixin):
+    """Quantized 3D maxpool"""
+    def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
+                 return_indices=False, ceil_mode=False, **kwargs):
+        super(QuantMaxPool3d, self).__init__(kernel_size, stride, padding, dilation,
+                                             return_indices, ceil_mode)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantMaxPool3d, self).forward(quant_input)
+
+
+class QuantAvgPool1d(pooling.AvgPool1d, _utils.QuantInputMixin):
+    """Quantized 1D average pool"""
+    def __init__(self, kernel_size, stride=None, padding=0, ceil_mode=False,
+                 count_include_pad=True, **kwargs):
+        super(QuantAvgPool1d, self).__init__(kernel_size, stride, padding, ceil_mode,
+                                             count_include_pad)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantAvgPool1d, self).forward(quant_input)
+
+class QuantAvgPool2d(pooling.AvgPool2d, _utils.QuantInputMixin):
+    """Quantized 2D average pool"""
+    def __init__(self, kernel_size, stride=None, padding=0, ceil_mode=False,
+                 count_include_pad=True, divisor_override=None, **kwargs):
+        super(QuantAvgPool2d, self).__init__(kernel_size, stride, padding, ceil_mode,
+                                             count_include_pad, divisor_override)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantAvgPool2d, self).forward(quant_input)
+
+
+class QuantAvgPool3d(pooling.AvgPool3d, _utils.QuantInputMixin):
+    """Quantized 3D average pool"""
+    def __init__(self, kernel_size, stride=None, padding=0, ceil_mode=False,
+                 count_include_pad=True, divisor_override=None, **kwargs):
+        super(QuantAvgPool3d, self).__init__(kernel_size, stride, padding, ceil_mode,
+                                             count_include_pad, divisor_override)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantAvgPool3d, self).forward(quant_input)
+
+
+class QuantAdaptiveAvgPool1d(pooling.AdaptiveAvgPool1d, _utils.QuantInputMixin):
+    """Quantized 1D adaptive average pool"""
+    def __init__(self, output_size, **kwargs):
+        super(QuantAdaptiveAvgPool1d, self).__init__(output_size)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantAdaptiveAvgPool1d, self).forward(quant_input)
+
+
+class QuantAdaptiveAvgPool2d(pooling.AdaptiveAvgPool2d, _utils.QuantInputMixin):
+    """Quantized 2D adaptive average pool"""
+    def __init__(self, output_size, **kwargs):
+        super(QuantAdaptiveAvgPool2d, self).__init__(output_size)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantAdaptiveAvgPool2d, self).forward(quant_input)
+
+
+class QuantAdaptiveAvgPool3d(pooling.AdaptiveAvgPool3d, _utils.QuantInputMixin):
+    """Quantized 3D adaptive average pool"""
+    def __init__(self, output_size, **kwargs):
+        super(QuantAdaptiveAvgPool3d, self).__init__(output_size)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantAdaptiveAvgPool3d, self).forward(quant_input)
+
+
+# Define alias with Quant prefix
+MaxPool1d = QuantMaxPool1d
+MaxPool2d = QuantMaxPool2d
+MaxPool3d = QuantMaxPool3d
+AvgPool1d = QuantAvgPool1d
+AvgPool2d = QuantAvgPool2d
+AvgPool3d = QuantAvgPool3d
+AdaptiveAvgPool1d = QuantAdaptiveAvgPool1d
+AdaptiveAvgPool2d = QuantAdaptiveAvgPool2d
+AdaptiveAvgPool3d = QuantAdaptiveAvgPool3d

tools/pytorch-quantization/pytorch_quantization/nn/modules/quant_rnn.py

@@ -0,0 +1,467 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+"""RNN implementation in python
+originally copied from https://github.com/pytorch/pytorch/blob/v0.4.1/torch/nn/modules/rnn.py
+    backend is changed to _functions/rnn.py
+"""
+import math
+import torch
+import warnings
+import itertools
+import numbers
+
+from torch import nn
+from torch.nn import Parameter
+from torch.nn.utils.rnn import PackedSequence
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.nn._functions import quant_rnn
+
+from . import _utils
+
+__all__ = ["QuantLSTM", "QuantLSTMCell", "LSTM", "LSTMCell"]
+
+class QuantRNNBase(nn.Module, _utils.QuantMixin):
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW
+
+    def __init__(self, mode, input_size, hidden_size,
+                 num_layers=1, bias=True, batch_first=False,
+                 dropout=0, bidirectional=False, proj_size=0, **kwargs):
+        super(QuantRNNBase, self).__init__()
+        self.mode = mode
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.bias = bias
+        self.batch_first = batch_first
+        self.dropout = dropout
+        self.dropout_state = {}
+        self.bidirectional = bidirectional
+        self.proj_size = proj_size
+        num_directions = 2 if bidirectional else 1
+
+        if not isinstance(dropout, numbers.Number) or not 0 <= dropout <= 1 or \
+                isinstance(dropout, bool):
+            raise ValueError("dropout should be a number in range [0, 1] "
+                             "representing the probability of an element being "
+                             "zeroed")
+        if dropout > 0 and num_layers == 1:
+            warnings.warn("dropout option adds dropout after all but last "
+                          "recurrent layer, so non-zero dropout expects "
+                          "num_layers greater than 1, but got dropout={} and "
+                          "num_layers={}".format(dropout, num_layers))
+
+        if proj_size < 0:
+            raise ValueError("proj_size should be a positive integer or zero to disable projections")
+        if proj_size > 0:
+            raise ValueError("proj_size is not supported in pytorch-quantization yet")
+
+        if mode == 'LSTM':
+            gate_size = 4 * hidden_size
+        elif mode == 'GRU':
+            gate_size = 3 * hidden_size
+        else:
+            gate_size = hidden_size
+
+        self._all_weights = []
+        for layer in range(num_layers):
+            for direction in range(num_directions):
+                layer_input_size = input_size if layer == 0 else hidden_size * num_directions
+
+                w_ih = Parameter(torch.Tensor(gate_size, layer_input_size))
+                w_hh = Parameter(torch.Tensor(gate_size, hidden_size))
+                b_ih = Parameter(torch.Tensor(gate_size))
+                b_hh = Parameter(torch.Tensor(gate_size))
+                layer_params = (w_ih, w_hh, b_ih, b_hh)
+
+                suffix = '_reverse' if direction == 1 else ''
+                param_names = ['weight_ih_l{}{}', 'weight_hh_l{}{}']
+                if bias:
+                    param_names += ['bias_ih_l{}{}', 'bias_hh_l{}{}']
+                param_names = [x.format(layer, suffix) for x in param_names]
+
+                for name, param in zip(param_names, layer_params):
+                    setattr(self, name, param)
+                self._all_weights.append(param_names)
+
+        self.flatten_parameters()
+        self.reset_parameters()
+
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        self.init_quantizer(quant_desc_input, quant_desc_weight, num_layers=num_layers * (1 + bidirectional))
+
+    def flatten_parameters(self):
+        """Resets parameter data pointer so that they can use faster code paths.
+
+        Right now, this works only if the module is on the GPU and cuDNN is enabled.
+        Otherwise, it's a no-op.
+        """
+        any_param = next(self.parameters()).data
+        if not any_param.is_cuda or not torch.backends.cudnn.is_acceptable(any_param):
+            self._data_ptrs = []
+            return
+
+        # If any parameters alias, we fall back to the slower, copying code path. This is
+        # a sufficient check, because overlapping parameter buffers that don't completely
+        # alias would break the assumptions of the uniqueness check in
+        # Module.named_parameters().
+        unique_data_ptrs = set(p.data_ptr() for l in self.all_weights for p in l)
+        if len(unique_data_ptrs) != sum(len(l) for l in self.all_weights):
+            self._data_ptrs = []
+            return
+
+        with torch.cuda.device_of(any_param):
+            import torch.backends.cudnn.rnn as rnn
+
+            weight_arr = list(itertools.chain.from_iterable(self.all_weights))
+            weight_stride0 = len(self.all_weights[0])
+
+            # NB: This is a temporary hack while we still don't have Tensor
+            # bindings for ATen functions
+            with torch.no_grad():
+                # NB: this is an INPLACE function on weight_arr, that's why the
+                # no_grad() is necessary.
+                weight_buf = torch._cudnn_rnn_flatten_weight(weight_arr, weight_stride0, self.input_size,
+                                                             rnn.get_cudnn_mode(self.mode), self.hidden_size,
+                                                             self.proj_size, self.num_layers, self.batch_first,
+                                                             bool(self.bidirectional))
+
+            self._param_buf_size = weight_buf.size(0)
+            self._data_ptrs = list(p.data.data_ptr() for p in self.parameters())
+
+    def _apply(self, fn):
+        ret = super(QuantRNNBase, self)._apply(fn)
+        self.flatten_parameters()
+        return ret
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.hidden_size)
+        for weight in self.parameters():
+            weight.data.uniform_(-stdv, stdv)
+
+    def check_forward_args(self, input, hidden, batch_sizes):
+        is_input_packed = batch_sizes is not None
+        expected_input_dim = 2 if is_input_packed else 3
+        if input.dim() != expected_input_dim:
+            raise RuntimeError(
+                'input must have {} dimensions, got {}'.format(
+                    expected_input_dim, input.dim()))
+        if self.input_size != input.size(-1):
+            raise RuntimeError(
+                'input.size(-1) must be equal to input_size. Expected {}, got {}'.format(
+                    self.input_size, input.size(-1)))
+
+        if is_input_packed:
+            mini_batch = int(batch_sizes[0])
+        else:
+            mini_batch = input.size(0) if self.batch_first else input.size(1)
+
+        num_directions = 2 if self.bidirectional else 1
+        expected_hidden_size = (self.num_layers * num_directions,
+                                mini_batch, self.hidden_size)
+
+        def check_hidden_size(hx, expected_hidden_size, msg='Expected hidden size {}, got {}'):
+            if tuple(hx.size()) != expected_hidden_size:
+                raise RuntimeError(msg.format(expected_hidden_size, tuple(hx.size())))
+
+        if self.mode == 'LSTM':
+            check_hidden_size(hidden[0], expected_hidden_size,
+                              'Expected hidden[0] size {}, got {}')
+            check_hidden_size(hidden[1], expected_hidden_size,
+                              'Expected hidden[1] size {}, got {}')
+        else:
+            check_hidden_size(hidden, expected_hidden_size)
+
+    def forward(self, input, hx=None):
+        is_packed = isinstance(input, PackedSequence)
+        if is_packed:
+            input, batch_sizes, sorted_indices, unsorted_indices = input
+            max_batch_size = batch_sizes[0]
+            max_batch_size = int(max_batch_size)
+        else:
+            batch_sizes = None
+            max_batch_size = input.size(0) if self.batch_first else input.size(1)
+
+        if hx is None:
+            num_directions = 2 if self.bidirectional else 1
+            hx = input.new_zeros(self.num_layers * num_directions,
+                                 max_batch_size, self.hidden_size,
+                                 requires_grad=False)
+            if self.mode == 'LSTM':
+                hx = (hx, hx)
+
+        has_flat_weights = list(p.data.data_ptr() for p in self.parameters()) == self._data_ptrs
+        if has_flat_weights:
+            first_data = next(self.parameters()).data
+            assert first_data.storage().size() == self._param_buf_size
+            flat_weight = first_data.new().set_(first_data.storage(), 0, torch.Size([self._param_buf_size]))
+        else:
+            flat_weight = None
+
+        self.check_forward_args(input, hx, batch_sizes)
+        func = quant_rnn.RNN(
+            self.mode,
+            self.input_size,
+            self.hidden_size,
+            num_layers=self.num_layers,
+            batch_first=self.batch_first,
+            dropout=self.dropout,
+            train=self.training,
+            bidirectional=self.bidirectional,
+            dropout_state=self.dropout_state,
+            variable_length=is_packed,
+            flat_weight=flat_weight
+        )
+        output, hidden = func(input, self.all_weights, hx, batch_sizes, self._input_quantizers, self._weight_quantizers)
+        if is_packed:
+            output = PackedSequence(output, batch_sizes)
+        return output, hidden
+
+    def extra_repr(self):
+        s = '{input_size}, {hidden_size}'
+        if self.num_layers != 1:
+            s += ', num_layers={num_layers}'
+        if self.bias is not True:
+            s += ', bias={bias}'
+        if self.batch_first is not False:
+            s += ', batch_first={batch_first}'
+        if self.dropout != 0:
+            s += ', dropout={dropout}'
+        if self.bidirectional is not False:
+            s += ', bidirectional={bidirectional}'
+        return s.format(**self.__dict__)
+
+    def __setstate__(self, d):
+        super(QuantRNNBase, self).__setstate__(d)
+        self.__dict__.setdefault('_data_ptrs', [])
+        if 'all_weights' in d:
+            self._all_weights = d['all_weights']
+        if isinstance(self._all_weights[0][0], str):
+            return
+        num_layers = self.num_layers
+        num_directions = 2 if self.bidirectional else 1
+        self._all_weights = []
+        for layer in range(num_layers):
+            for direction in range(num_directions):
+                suffix = '_reverse' if direction == 1 else ''
+                weights = ['weight_ih_l{}{}', 'weight_hh_l{}{}', 'bias_ih_l{}{}', 'bias_hh_l{}{}']
+                weights = [x.format(layer, suffix) for x in weights]
+                if self.bias:
+                    self._all_weights += [weights]
+                else:
+                    self._all_weights += [weights[:2]]
+
+    @property
+    def all_weights(self):
+        return [[getattr(self, weight) for weight in weights] for weights in self._all_weights]
+
+
+class QuantRNN(QuantRNNBase):
+    r"""Applies a multi-layer Elman RNN with `tanh` or `ReLU` non-linearity to an
+    input sequence.
+    """
+
+    def __init__(self, *args, **kwargs):
+        if 'proj_size' in kwargs:
+            raise ValueError("proj_size argument is only supported for LSTM, not RNN or GRU")
+        if 'nonlinearity' in kwargs:
+            if kwargs['nonlinearity'] == 'tanh':
+                mode = 'RNN_TANH'
+            elif kwargs['nonlinearity'] == 'relu':
+                mode = 'RNN_RELU'
+            else:
+                raise ValueError("Unknown nonlinearity '{}'".format(
+                    kwargs['nonlinearity']))
+            del kwargs['nonlinearity']
+        else:
+            mode = 'RNN_TANH'
+
+        super(QuantRNN, self).__init__(mode, *args, **kwargs)
+
+
+class QuantLSTM(QuantRNNBase):
+    r"""Applies a multi-layer long short-term memory (LSTM) RNN to an input
+    sequence.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(QuantLSTM, self).__init__('LSTM', *args, **kwargs)
+
+
+class GRU(QuantRNNBase):
+    r"""Applies a multi-layer gated recurrent unit (GRU) RNN to an input sequence.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(GRU, self).__init__('GRU', *args, **kwargs)
+
+class QuantRNNCellBase(nn.Module, _utils.QuantMixin):
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW
+
+    def extra_repr(self):
+        s = '{input_size}, {hidden_size}'
+        if 'bias' in self.__dict__ and self.bias is not True:
+            s += ', bias={bias}'
+        if 'nonlinearity' in self.__dict__ and self.nonlinearity != "tanh":
+            s += ', nonlinearity={nonlinearity}'
+        return s.format(**self.__dict__)
+
+    def check_forward_input(self, input):
+        if input.size(1) != self.input_size:
+            raise RuntimeError(
+                "input has inconsistent input_size: got {}, expected {}".format(
+                    input.size(1), self.input_size))
+
+    def check_forward_hidden(self, input, hx, hidden_label=''):
+        if input.size(0) != hx.size(0):
+            raise RuntimeError(
+                "Input batch size {} doesn't match hidden{} batch size {}".format(
+                    input.size(0), hidden_label, hx.size(0)))
+
+        if hx.size(1) != self.hidden_size:
+            raise RuntimeError(
+                "hidden{} has inconsistent hidden_size: got {}, expected {}".format(
+                    hidden_label, hx.size(1), self.hidden_size))
+
+
+class QuantRNNCell(QuantRNNCellBase):
+    r"""An Elman RNN cell with tanh or ReLU non-linearity.
+    """
+
+    def __init__(self, input_size, hidden_size, bias=True, nonlinearity="tanh"):
+        super(QuantRNNCell, self).__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.bias = bias
+        self.nonlinearity = nonlinearity
+        self.weight_ih = Parameter(torch.Tensor(hidden_size, input_size))
+        self.weight_hh = Parameter(torch.Tensor(hidden_size, hidden_size))
+        if bias:
+            self.bias_ih = Parameter(torch.Tensor(hidden_size))
+            self.bias_hh = Parameter(torch.Tensor(hidden_size))
+        else:
+            self.register_parameter('bias_ih', None)
+            self.register_parameter('bias_hh', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.hidden_size)
+        for weight in self.parameters():
+            weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, hx=None):
+        self.check_forward_input(input)
+        if hx is None:
+            hx = input.new_zeros(input.size(0), self.hidden_size, requires_grad=False)
+        self.check_forward_hidden(input, hx)
+        if self.nonlinearity == "tanh":
+            func = quant_rnn.RNNTanhCell
+        elif self.nonlinearity == "relu":
+            func = quant_rnn.RNNReLUCell
+        else:
+            raise RuntimeError(
+                "Unknown nonlinearity: {}".format(self.nonlinearity))
+
+        return func(
+            input, hx,
+            self.weight_ih, self.weight_hh,
+            self.bias_ih, self.bias_hh,
+        )
+
+
+class QuantLSTMCell(QuantRNNCellBase):
+    r"""A long short-term memory (LSTM) cell.
+    """
+
+    def __init__(self, input_size, hidden_size, bias=True, **kwargs):
+        super(QuantLSTMCell, self).__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.bias = bias
+        self.weight_ih = Parameter(torch.Tensor(4 * hidden_size, input_size))
+        self.weight_hh = Parameter(torch.Tensor(4 * hidden_size, hidden_size))
+        if bias:
+            self.bias_ih = Parameter(torch.Tensor(4 * hidden_size))
+            self.bias_hh = Parameter(torch.Tensor(4 * hidden_size))
+        else:
+            self.register_parameter('bias_ih', None)
+            self.register_parameter('bias_hh', None)
+        self.reset_parameters()
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        self.init_quantizer(quant_desc_input, quant_desc_weight)
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.hidden_size)
+        for weight in self.parameters():
+            weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, hx=None):
+        self.check_forward_input(input)
+        if hx is None:
+            hx = input.new_zeros(input.size(0), self.hidden_size, requires_grad=False)
+            hx = (hx, hx)
+        self.check_forward_hidden(input, hx[0], '[0]')
+        self.check_forward_hidden(input, hx[1], '[1]')
+        return quant_rnn.LSTMCell(
+            input, hx,
+            self.weight_ih, self.weight_hh,
+            self.bias_ih, self.bias_hh,
+            self._input_quantizer, self._weight_quantizer
+        )
+
+
+class GRUCell(QuantRNNCellBase):
+    r"""A gated recurrent unit (GRU) cell
+    """
+
+    def __init__(self, input_size, hidden_size, bias=True):
+        super(GRUCell, self).__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.bias = bias
+        self.weight_ih = Parameter(torch.Tensor(3 * hidden_size, input_size))
+        self.weight_hh = Parameter(torch.Tensor(3 * hidden_size, hidden_size))
+        if bias:
+            self.bias_ih = Parameter(torch.Tensor(3 * hidden_size))
+            self.bias_hh = Parameter(torch.Tensor(3 * hidden_size))
+        else:
+            self.register_parameter('bias_ih', None)
+            self.register_parameter('bias_hh', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.hidden_size)
+        for weight in self.parameters():
+            weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, hx=None):
+        self.check_forward_input(input)
+        if hx is None:
+            hx = input.new_zeros(input.size(0), self.hidden_size, requires_grad=False)
+        self.check_forward_hidden(input, hx)
+        return quant_rnn.GRUCell(
+            input, hx,
+            self.weight_ih, self.weight_hh,
+            self.bias_ih, self.bias_hh,
+        )
+
+LSTM = QuantLSTM
+LSTMCell = QuantLSTMCell

tools/pytorch-quantization/pytorch_quantization/nn/modules/tensor_quantizer.py

@@ -0,0 +1,456 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""TensorQuantizer Module"""
+import math
+from absl import logging
+
+import torch
+from torch import nn
+
+from pytorch_quantization.tensor_quant import QuantDescriptor, tensor_quant, fake_tensor_quant, scaled_e4m3
+from pytorch_quantization.nn.modules.clip import Clip
+
+from pytorch_quantization import calib
+
+import pytorch_quantization.utils as quant_utils
+
+__all__ = ['TensorQuantizer']
+
+
+class TensorQuantizer(nn.Module):
+    """Tensor quantizer module
+
+    This module uses tensor_quant or fake_tensor_quant function to quantize a tensor. And wrappers variable, moving
+    statistics we'd want when training a quantized network.
+
+    Experimental features:
+        ``clip`` stage learns range before enabling quantization.
+        ``calib`` stage runs calibration
+
+    Args:
+        quant_desc: An instance of :func:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+        disabled: A boolean. If True, by pass the whole module returns input. Default False.
+        if_quant: A boolean. If True, run main quantization body. Default True.
+        if_clip: A boolean. If True, clip before quantization and learn amax. Default False.
+        if_calib: A boolean. If True, run calibration. Not implemented yet. Settings of calibration will probably
+            go to :func:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+
+    Raises:
+
+    Readonly Properties:
+        - axis:
+        - fake_quant:
+        - scale:
+        - step_size:
+
+    Mutable Properties:
+        - num_bits:
+        - unsigned:
+        - amax:
+    """
+
+    _enable_onnx_export = False
+
+    def __init__(self, quant_desc=QuantDescriptor(), disabled=False, if_quant=True, if_clip=False, if_calib=False):
+        """Initialize quantizer and set up required variables"""
+        super(TensorQuantizer, self).__init__()
+        # Expand quant_desc. Use quant_desc.dict would be eaiser, but adding one-by-one explicitly gives more control
+        self._num_bits = quant_desc.num_bits
+        self._fake_quant = quant_desc.fake_quant
+        self._axis = quant_desc.axis
+        self._scale_amax = quant_desc.scale_amax
+        self._learn_amax = quant_desc.learn_amax
+        self._unsigned = quant_desc.unsigned
+        self._narrow_range = quant_desc.narrow_range
+
+        self._scale = None if not quant_desc.fake_quant else 1.
+        self._disabled = disabled
+        self._if_quant = if_quant
+        self._if_clip = False
+        self._if_calib = if_calib
+
+        if quant_desc.amax is not None:
+            self.register_buffer('_amax', torch.tensor(quant_desc.amax))
+
+        # Clip module consumes a lot of memory, so only create it if learn_amax is True
+        if self._learn_amax:
+            init_amax = quant_desc.amax if quant_desc.amax is not None else 1.
+            self.clip = Clip(-init_amax, init_amax, learn_min=True, learn_max=True)
+            # It makes more sense to enable clip stage (which learns amax) if learn_amax is true
+            self.enable_clip()
+        if if_clip:
+            self.enable_clip()
+
+        if quant_desc.calib_method == "histogram":
+            logging.info("Creating histogram calibrator")
+            self._calibrator = calib.HistogramCalibrator(num_bits=self._num_bits,
+                                                         axis=self._axis,
+                                                         unsigned=self._unsigned)
+        elif quant_desc.calib_method == "max":
+            logging.info("Creating Max calibrator")
+            self._calibrator = calib.MaxCalibrator(num_bits=self._num_bits, axis=self._axis, unsigned=self._unsigned)
+
+    # pylint:disable=missing-docstring
+    @property
+    def num_bits(self):
+        return self._num_bits
+
+    @property
+    def maxbound(self):
+        if self._num_bits == (4, 3):
+            return 448.0
+        return (1 << (self._num_bits - 1 + int(self._unsigned))) - 1
+
+    @property
+    def unsigned(self):
+        return self._unsigned
+
+    @property
+    def scale(self):
+        if self._fake_quant:
+            logging.error("Fake quantize mode doesn't use scale explicitly!")
+        if self._scale is None:
+            logging.critical("Accessing scale before quantizing any tensor!")
+        return self._scale
+
+    @property
+    def pre_quant_scale(self):
+        if not hasattr(self, "_pre_quant_scale"):
+            return None
+        return self._pre_quant_scale
+
+    @property
+    def amax(self):
+        if not hasattr(self, "_amax"):
+            return None
+        return self._amax
+
+    @property
+    def step_size(self):
+        if not hasattr(self, "_amax"):
+            logging.error("step_size is undefined under dynamic amax mode!")
+            return None
+        return self._amax / (2.0**(self._num_bits - 1 + int(self._unsigned)) - 1.0)
+
+    @property
+    def axis(self):
+        return self._axis
+
+    @property
+    def fake_quant(self):
+        return self._fake_quant
+
+    @property
+    def narrow_range(self):
+        return self._narrow_range
+
+    def disable(self):
+        """Bypass the module"""
+        self._disabled = True
+
+    def enable(self):
+        self._disabled = False
+
+    def disable_clip(self):
+        """Disable clip stage"""
+        self._if_clip = False
+        self.clip.clip_value_min.requires_grad = False
+        self.clip.clip_value_max.requires_grad = False
+
+    def enable_clip(self):
+        """Enable clip stage"""
+        logging.warning("Enable `clip` stage for amax learning.")
+        if not self._learn_amax:
+            raise ValueError("learn_amax is False. Cannot enable clip.")
+        self.clip.clip_value_min.requires_grad = True
+        self.clip.clip_value_max.requires_grad = True
+        self._if_clip = True
+
+    def disable_calib(self):
+        logging.warning("Disable {}".format(self._calibrator.__class__.__name__))
+        self._if_calib = False
+
+    def enable_calib(self):
+        if self._calibrator is None:
+            raise ValueError("Calibrator was not created, cannot enable calibration.")
+        logging.info("Enable {}".format(self._calibrator.__class__.__name__))
+        self._if_calib = True
+
+    def disable_quant(self):
+        logging.info("Disable `quant` stage.")
+        self._if_quant = False
+
+    def enable_quant(self):
+        logging.info("Enable `quant` stage.")
+        self._if_quant = True
+
+    @amax.setter
+    def amax(self, value):
+        if value is None:
+            logging.error("Setting amax no None is meaningless.")
+        else:
+            if isinstance(value, torch.Tensor):
+                logging.warning("amax setter is not designed to take tensor.")
+            if not hasattr(self, "_amax"):
+                self.register_buffer('_amax', torch.tensor(value))
+            else:
+                value = torch.tensor(value, device=self._amax.device)
+                if self._amax.shape != value.shape:
+                    raise RuntimeError("Changing shape when setting amax is not allowed.")
+                self._amax.data.copy_(value.data)
+
+    @pre_quant_scale.setter
+    def pre_quant_scale(self, value):
+        if value is None:
+            logging.error("Setting pre_quant_scale no None is meaningless.")
+        else:
+            if not hasattr(self, "_pre_quant_scale"):
+                self.register_buffer('_pre_quant_scale', torch.tensor(value))
+            else:
+                value = torch.tensor(value, device=self._pre_quant_scale.device)
+                if self._pre_quant_scale.shape != value.shape:
+                    raise RuntimeError("Changing shape when setting pre_quant_scale is not allowed.")
+                self._pre_quant_scale.data.copy_(value.data)
+
+    @num_bits.setter
+    def num_bits(self, value):
+        self._num_bits = value
+
+    @unsigned.setter
+    def unsigned(self, value):
+        self._unsigned = value
+
+    @narrow_range.setter
+    def narrow_range(self, value):
+        self._narrow_range = value
+
+    # pylint:enable=missing-docstring
+    def load_calib_amax(self, *args, **kwargs):
+        """Load amax from calibrator.
+
+        Updates the amax buffer with value computed by the calibrator, creating it if necessary.
+        *args and **kwargs are directly passed to compute_amax, except "strict" in kwargs. Refer to
+        compute_amax for more details.
+        """
+        strict = kwargs.pop("strict", True)
+        if getattr(self, '_calibrator', None) is None:
+            raise RuntimeError("Calibrator not created.")
+        calib_amax = self._calibrator.compute_amax(*args, **kwargs)
+        if calib_amax is None:
+            err_msg = "Calibrator returned None. This usually happens when calibrator hasn't seen any tensor."
+            if not strict:
+                logging.warning(err_msg)
+                logging.warning("Set amax to NaN!")
+                calib_amax = torch.tensor(math.nan)
+            else:
+                raise RuntimeError(err_msg + " Passing 'strict=False' to `load_calib_amax()` will ignore the error.")
+        logging.warning("Load calibrated amax, shape={}.".format(calib_amax.shape))
+        logging.log_first_n(logging.WARNING, "Call .cuda() if running on GPU after loading calibrated amax.", 1)
+        if not hasattr(self, '_amax'):
+            self.register_buffer("_amax", calib_amax.data)
+        else:
+            self._amax.copy_(calib_amax)
+
+    def init_learn_amax(self):
+        """Initialize learned amax from fixed amax"""
+        if self._learn_amax is False:
+            raise RuntimeError("Called init_learn_amax with learn_amax=False.")
+        logging.warning("Load amax as initial value for amax learning!")
+        if self._amax.numel() != 1:
+            logging.warning("Per channel learned amax not supported. Initializing with max(amax).")
+            init_amax = torch.max(self._amax)
+        else:
+            init_amax = self._amax
+        self.clip.clip_value_min.data.copy_(-init_amax.data)
+        self.clip.clip_value_max.data.copy_(init_amax.data)
+
+    def _get_amax(self, inputs):
+        """get amax from buffer or compute it dynamically."""
+        if hasattr(self, '_amax'):
+            amax = self._amax
+        else:
+            if self._axis is None:
+                reduce_axis = None
+            else:
+                reduce_axis = []
+                # Swap axis to reduce
+                axis = self._axis if isinstance(self._axis, (list, tuple)) else [self._axis]
+                for i in range(inputs.dim()):
+                    if not i in axis:
+                        reduce_axis.append(i)
+            amax = quant_utils.reduce_amax(inputs, axis=reduce_axis, keepdims=True).detach()
+        if self._scale_amax is not None:
+            amax = amax.detach() * self._scale_amax
+
+        amax = amax.data
+
+        # cast amax to float32 if it is in a lower precision dtype
+        if amax.dtype not in (torch.double, torch.float):
+            amax = amax.float()
+
+        return amax
+
+    def _quant_forward(self, inputs):
+        """Quantized forward pass."""
+        if self._learn_amax:
+            inputs = self.clip(inputs)
+            amax = torch.max(-self.clip.clip_value_min, self.clip.clip_value_max).detach()
+        else:
+            amax = self._get_amax(inputs)
+
+        if self._fake_quant:
+            outputs = fake_tensor_quant(inputs, amax, self._num_bits, self._unsigned, self._narrow_range)
+        else:
+            outputs, self._scale = tensor_quant(inputs, amax, self._num_bits, self._unsigned)
+
+        return outputs
+
+    def _check_onnx_readiness(self, inputs):
+        """Check if quantizer is ready for ONNX export."""
+
+        assert hasattr(
+            self, '_amax'), ("Quantizer has not been calibrated. ONNX export requires the quantizer to be calibrated."
+                             "Calibrate and load amax before exporting to ONNX.")
+
+        if self._if_calib:
+            logging.warning("Quantizer is in calibration mode. "
+                            "Please complete calibration before exporting to ONNX for correct results.")
+
+        amax = self._get_amax(inputs)
+
+        # We only support scalar amax for E4M3 ONNX export
+        if isinstance(self.num_bits, tuple):
+            assert amax.numel() == 1, ("E4M3 supports ONNX export only for per-tensor quantization."
+                                       " Per-tensor quantization requires scalar amax. "
+                                       f"Received non-scalar amax of shape: {amax.shape}")
+
+    def forward(self, inputs):
+        """Apply tensor_quant function to inputs
+
+        Args:
+            inputs: A Tensor of type float32.
+
+        Returns:
+            outputs: A Tensor of type output_dtype
+        """
+
+        if self._enable_onnx_export:
+            self._check_onnx_readiness(inputs)
+
+        # Activation scaling for smoothquant
+        if self.pre_quant_scale is not None:
+            inputs = inputs * self.pre_quant_scale
+
+        if self._disabled:
+            return inputs
+
+        outputs = inputs
+
+        if self._if_calib:
+            if self._calibrator is None:
+                raise RuntimeError("Calibrator was not created.")
+            # Shape is only known when it sees the first tensor
+            self._calibrator.collect(inputs)
+
+        if self._if_clip:
+            if not self._learn_amax:
+                raise RuntimeError("Clip without learning amax is not implemented.")
+            outputs = self.clip(inputs)
+
+        if self._if_quant:
+            if not isinstance(self._num_bits, tuple):
+                outputs = self._quant_forward(inputs)
+            else:
+                E, M = self._num_bits
+                outputs = scaled_e4m3(inputs, self._get_amax(inputs), E, M)
+
+        return outputs
+
+    def _short_amax(self, fmt='.4f'):
+        """Short description of amax
+
+        Returns:
+            'dynamic': if _amax is not registered
+            'amax': if _amax is per-tensor
+            '[min, max](size)': if _amax is per-channel
+        """
+        if not hasattr(self, '_amax'):
+            return 'dynamic'
+        if self._amax is None:
+            return "None"
+        if self._amax.numel() == 1:
+            return '{:{fmt}}'.format(self._amax.item(), fmt=fmt)
+        return '[{:{fmt}}, {:{fmt}}]({})'.format(self._amax.min().item(),
+                                                 self._amax.max().item(),
+                                                 self._amax.numel(),
+                                                 fmt=fmt)
+
+    def extra_repr(self):
+        if self._disabled:
+            return "disabled"
+        s = "{}{}bit".format("unsigned " if self._unsigned else "", self._num_bits)
+        s += " narrow" if (self._narrow_range) else ""
+        s += " fake" if (self._fake_quant) else ""
+        s += " axis={}".format(self._axis) if self._axis is not None else " per-tensor"
+        s += " amax={}".format(self._short_amax())
+        s += " *{}".format(self._scale_amax) if self._scale_amax else ""
+        s += " pre_quant_scale" if self.pre_quant_scale is not None else ""
+        s += " learned" if (self._learn_amax) else ""
+        s += " calibrator={}".format(self._calibrator.__class__.__name__) if (self._calibrator is not None) else ""
+        s += " scale={}".format(self._scale) if self._scale is not None else ""
+        s += " quant" if (self._if_quant) else ""
+        s += " clip" if (self._if_clip) else ""
+        s += " calib" if (self._if_calib) else ""
+        return s
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        """Overloaded module function
+
+        Adds warnings during state_dict loading.
+        A workaround is implemented for loading amax from checkpoint and only supports CUDA.
+
+        Args:
+            state_dict: A dict containing the state of the top level module
+            prefix: A string that prefixes all of this modules state in state_dict, e.g. 'model.conv1.'
+        """
+        dst_has_amax = '_amax' in self._buffers
+        src_has_amax = prefix + '_amax' in state_dict
+
+        if not src_has_amax and dst_has_amax:
+            logging.error("{}: No amax in state_dict.".format(prefix[:-1]))
+        elif src_has_amax and not dst_has_amax:
+            logging.debug(("{}: No '_amax' buffer to load amax into."
+                           " '_amax` will be created as WAR for now. "
+                           "This behavior will change in future.").format(prefix[:-1]))
+            self.register_buffer("_amax", state_dict[prefix + '_amax'].data.cuda())
+        elif src_has_amax and dst_has_amax:
+            logging.warning("{}: Overwriting amax.".format(prefix[:-1]))
+
+        dst_has_pre_quant_scale = '_pre_quant_scale' in self._buffers
+        src_has_pre_quant_scale = prefix + '_pre_quant_scale' in state_dict
+
+        if not src_has_pre_quant_scale and dst_has_pre_quant_scale:
+            logging.error("{}: No pre_quant_scale in state_dict.".format(prefix[:-1]))
+        elif src_has_pre_quant_scale and not dst_has_pre_quant_scale:
+            logging.debug(("{}: No '_pre_quant_scale' buffer to load pre_quant_scale into."
+                           " '_pre_quant_scale` will be created as WAR for now. "
+                           "This behavior will change in future.").format(prefix[:-1]))
+            self.register_buffer("_pre_quant_scale", state_dict[prefix + '_pre_quant_scale'].data.cuda())
+        elif src_has_pre_quant_scale and dst_has_pre_quant_scale:
+            logging.warning("{}: Overwriting pre_quant_scale.".format(prefix[:-1]))
+
+        super(TensorQuantizer, self)._load_from_state_dict(state_dict, prefix, *args, **kwargs)