chore: import upstream snapshot with attribution

tools/pytorch-quantization/pytorch_quantization/__init__.py

@@ -0,0 +1,22 @@
+#
+# 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.
+#
+
+from absl import logging
+from .version import __version__
+from .quant_modules import *
+
+logging.use_absl_handler()

tools/pytorch-quantization/pytorch_quantization/calib/__init__.py

@@ -0,0 +1,24 @@
+#
+# 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.
+#
+
+
+"""``pytorch_quantization.calib`` provides Calibrator classes that
+collect data statistics and determine pytorch_quantization parameters.
+"""
+
+from .max import MaxCalibrator
+from .histogram import *

tools/pytorch-quantization/pytorch_quantization/calib/calibrator.py

@@ -0,0 +1,61 @@
+#
+# 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.
+#
+
+
+"""Abstract base class for calibrators"""
+
+
+class _Calibrator():
+    """Abstract base class of calibrators
+    Args:
+        num_bits: An integer. Number of bits of quantization.
+        axis: A tuple. see QuantDescriptor.
+        unsigned: A boolean. using unsigned quantization.
+
+    Readonly Properties:
+        axis:
+    """
+    def __init__(self, num_bits, axis, unsigned):
+        self._num_bits = num_bits
+        self._axis = axis
+        self._unsigned = unsigned
+
+    def collect(self, x):
+        """Abstract method: collect tensor statistics used to compute amax
+
+        Args:
+            x: A tensor
+        """
+        raise NotImplementedError
+
+    def reset(self):
+        """Abstract method: reset calibrator to initial state"""
+        raise NotImplementedError
+
+    def compute_amax(self, *args, **kwargs):
+        """Abstract method: compute the amax from the collected data
+
+        Returns:
+            amax: a tensor
+        """
+        raise NotImplementedError
+
+    def __repr__(self):
+        s = "num_bits={_num_bits}"
+        s += " axis={_axis}"
+        s += " unsigned={_unsigned}"
+        return s.format(**self.__dict__)

tools/pytorch-quantization/pytorch_quantization/calib/histogram.py

@@ -0,0 +1,386 @@
+#
+# 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.
+#
+"""Histogram based calibrators"""
+from collections import Counter
+import numpy as np
+from scipy.stats import entropy
+
+from absl import logging
+
+import torch
+
+from pytorch_quantization.calib.calibrator import _Calibrator
+from pytorch_quantization.tensor_quant import fake_tensor_quant, scaled_e4m3
+from pytorch_quantization import nn as quant_nn
+from pytorch_quantization import utils as quant_utils
+
+__all__ = ["HistogramCalibrator", "calibrate_weights"]
+
+
+class HistogramCalibrator(_Calibrator):
+    """Unified histogram calibrator
+
+    Histogram will be only collected once. compute_amax() performs entropy, percentile, or mse
+        calibration based on arguments
+
+    Args:
+        num_bits: An integer. Number of bits of quantization.
+        axis: A tuple. see QuantDescriptor.
+        unsigned: A boolean. using unsigned quantization.
+        num_bins: An integer. Number of histograms bins. Default 2048.
+        grow_method: A string. DEPRECATED. default None.
+        skip_zeros: A boolean. If True, skips zeros when collecting data for histogram. Default False.
+        torch_hist: A boolean. If True, collect histogram by torch.histc instead of np.histogram. If input tensor
+            is on GPU, histc will also be running on GPU. Default True.
+    """
+
+    def __init__(self, num_bits, axis, unsigned, num_bins=2048, grow_method=None, skip_zeros=False, torch_hist=True):
+        super(HistogramCalibrator, self).__init__(num_bits, axis, unsigned)
+        self._num_bins = num_bins
+        self._skip_zeros = skip_zeros
+
+        self._calib_bin_edges = None
+        self._calib_hist = None
+
+        self._torch_hist = torch_hist
+
+        if axis is not None:
+            raise NotImplementedError("Calibrator histogram collection only supports per tensor scaling")
+
+        if grow_method is not None:
+            logging.warning("grow_method is deprecated. Got %s, ingored!", grow_method)
+
+    def collect(self, x):
+        """Collect histogram"""
+        if torch.min(x) < 0.:
+            logging.log_first_n(logging.INFO,
+                                ("Calibrator encountered negative values. It shouldn't happen after ReLU. "
+                                 "Make sure this is the right tensor to calibrate."), 1)
+            x = x.abs()
+
+        x = x.float()
+
+        if not self._torch_hist:
+            x_np = x.cpu().detach().numpy()
+
+            if self._skip_zeros:
+                x_np = x_np[np.where(x_np != 0)]
+
+            if self._calib_bin_edges is None and self._calib_hist is None:
+                # first time it uses num_bins to compute histogram.
+                self._calib_hist, self._calib_bin_edges = np.histogram(x_np, bins=self._num_bins)
+            else:
+                temp_amax = np.max(x_np)
+                if temp_amax > self._calib_bin_edges[-1]:
+                    # increase the number of bins
+                    width = self._calib_bin_edges[1] - self._calib_bin_edges[0]
+                    # NOTE: np.arange may create an extra bin after the one containing temp_amax
+                    new_bin_edges = np.arange(self._calib_bin_edges[-1] + width, temp_amax + width, width)
+                    self._calib_bin_edges = np.hstack((self._calib_bin_edges, new_bin_edges))
+                hist, self._calib_bin_edges = np.histogram(x_np, bins=self._calib_bin_edges)
+                hist[:len(self._calib_hist)] += self._calib_hist
+                self._calib_hist = hist
+        else:
+            # This branch of code is designed to match numpy version as close as possible
+            with torch.no_grad():
+                if self._skip_zeros:
+                    x = x[torch.where(x != 0)]
+
+                # Because we collect histogram on absolute value, setting min=0 simplifying the rare case where
+                # minimum value is not exactly 0 and first batch collected has larger min value than later batches
+                x_max = x.max()
+                if self._calib_bin_edges is None and self._calib_hist is None:
+                    self._calib_hist = torch.histc(x, bins=self._num_bins, min=0, max=x_max)
+                    self._calib_bin_edges = torch.linspace(0, x_max, self._num_bins + 1)
+                else:
+                    if x_max > self._calib_bin_edges[-1]:
+                        width = self._calib_bin_edges[1] - self._calib_bin_edges[0]
+                        self._num_bins = int((x_max / width).ceil().item())
+                        self._calib_bin_edges = torch.arange(0, x_max + width, width, device=x.device)
+
+                    hist = torch.histc(x, bins=self._num_bins, min=0, max=self._calib_bin_edges[-1])
+                    hist[:self._calib_hist.numel()] += self._calib_hist
+                    self._calib_hist = hist
+
+    def reset(self):
+        """Reset the collected histogram"""
+        self._calib_bin_edges = None
+        self._calib_hist = None
+
+    def compute_amax(self, method: str, *, stride: int = 1, start_bin: int = 128, percentile: float = 99.99):
+        """Compute the amax from the collected histogram
+
+        Args:
+            method: A string. One of ['entropy', 'mse', 'percentile']
+
+        Keyword Arguments:
+            stride: An integer. Default 1
+            start_bin: An integer. Default 128
+            percentils: A float number between [0, 100]. Default 99.99.
+
+        Returns:
+            amax: a tensor
+        """
+        if isinstance(self._calib_hist, torch.Tensor):
+            calib_hist = self._calib_hist.to(torch.int64).cpu().numpy()
+            calib_bin_edges = self._calib_bin_edges.cpu().numpy()
+        else:
+            calib_hist = self._calib_hist
+            calib_bin_edges = self._calib_bin_edges
+
+        if method == 'entropy':
+            calib_amax = _compute_amax_entropy(calib_hist, calib_bin_edges, self._num_bits, self._unsigned, stride,
+                                               start_bin)
+        elif method == 'mse':
+            calib_amax = _compute_amax_mse(calib_hist, calib_bin_edges, self._num_bits, self._unsigned, stride,
+                                           start_bin)
+        elif method == 'percentile':
+            calib_amax = _compute_amax_percentile(calib_hist, calib_bin_edges, percentile)
+        else:
+            raise TypeError("Unknown calibration method {}".format(method))
+
+        return calib_amax
+
+    # pylint:disable=missing-docstring
+    def __str__(self):
+        s = "HistogramCalibrator("
+        if self._calib_bin_edges is None:
+            bin_edge_str = "None"
+        else:
+            bin_edge_str = "[{:.3f}, ..., {:.3f}]({})".format(self._calib_bin_edges[0], self._calib_bin_edges[-1],
+                                                              len(self._calib_bin_edges))
+        s += "calib_bin_edges={})".format(bin_edge_str)
+        return s
+
+    def __repr__(self):
+        s = "HistogramCalibrator("
+        s += super(HistogramCalibrator, self).__repr__()
+        s += " calib_bin_edges={_calib_bin_edges}"
+        s += " calib_hist={_calib_hist})"
+        return s.format(**self.__dict__)
+
+    # pylint:enable=missing-docstring
+
+
+# Ideally, we want to decouple collector (collect histogram) and calibrator (compute amax) as opposed to
+# the current calibrator design. The following compute amax functions are broken out from the calibrator
+# as first step towards there.
+def _compute_amax_entropy(calib_hist, calib_bin_edges, num_bits, unsigned, stride=1, start_bin=128):
+    """Returns amax that minimizes KL-Divergence of the collected histogram"""
+
+    # If calibrator hasn't collected any data, return none
+    if calib_bin_edges is None and calib_hist is None:
+        return None
+
+    def _normalize_distr(distr):
+        summ = np.sum(distr)
+        if summ != 0:
+            distr = distr / summ
+
+    bins = calib_hist[:]
+    bins[0] = bins[1]
+
+    total_data = np.sum(bins)
+
+    divergences = []
+    arguments = []
+
+    # we are quantizing to 128 values + sign if num_bits=8
+    nbins = 1 << (num_bits - 1 + int(unsigned))
+
+    starting = start_bin
+    stop = len(bins)
+
+    new_density_counts = np.zeros(nbins, dtype=np.float64)
+
+    for i in range(starting, stop + 1, stride):
+        new_density_counts.fill(0)
+        space = np.linspace(0, i, num=nbins + 1)
+        digitized_space = np.digitize(range(i), space) - 1
+
+        digitized_space[bins[:i] == 0] = -1
+
+        for idx, digitized in enumerate(digitized_space):
+            if digitized != -1:
+                new_density_counts[digitized] += bins[idx]
+
+        counter = Counter(digitized_space)
+        for key, val in counter.items():
+            if key != -1:
+                new_density_counts[key] = new_density_counts[key] / val
+
+        new_density = np.zeros(i, dtype=np.float64)
+        for idx, digitized in enumerate(digitized_space):
+            if digitized != -1:
+                new_density[idx] = new_density_counts[digitized]
+
+        total_counts_new = np.sum(new_density) + np.sum(bins[i:])
+        _normalize_distr(new_density)
+
+        reference_density = np.array(bins[:len(digitized_space)])
+        reference_density[-1] += np.sum(bins[i:])
+
+        total_counts_old = np.sum(reference_density)
+        if round(total_counts_new) != total_data or round(total_counts_old) != total_data:
+            raise RuntimeError("Count mismatch! total_counts_new={}, total_counts_old={}, total_data={}".format(
+                total_counts_new, total_counts_old, total_data))
+
+        _normalize_distr(reference_density)
+
+        ent = entropy(reference_density, new_density)
+        divergences.append(ent)
+        arguments.append(i)
+
+    divergences = np.array(divergences)
+    logging.debug("divergences={}".format(divergences))
+    last_argmin = len(divergences) - 1 - np.argmin(divergences[::-1])
+    calib_amax = calib_bin_edges[last_argmin * stride + starting]
+    calib_amax = torch.tensor(calib_amax.item())  #pylint: disable=not-callable
+
+    return calib_amax
+
+
+def _compute_amax_mse(calib_hist, calib_bin_edges, num_bits, unsigned, stride=1, start_bin=128):
+    """Returns amax that minimizes MSE of the collected histogram"""
+
+    # If calibrator hasn't collected any data, return none
+    if calib_bin_edges is None and calib_hist is None:
+        return None
+
+    counts = torch.from_numpy(calib_hist[:]).float().cuda()
+    edges = torch.from_numpy(calib_bin_edges[:]).float().cuda()
+    centers = (edges[1:] + edges[:-1]) / 2
+
+    mses = []
+    arguments = []
+
+    for i in range(start_bin, len(centers), stride):
+
+        amax = centers[i]
+        if isinstance(num_bits, int) and num_bits >= 0:
+            if num_bits == 0:
+                logging.error("num_bits is 0. This will result in the tensor being quantized to all zeros."
+                              " This mode should only be used for debugging purposes.")
+            quant_centers = fake_tensor_quant(centers, amax, num_bits, unsigned)
+        elif num_bits == (4, 3):
+            quant_centers = scaled_e4m3(centers, amax, num_bits[0], num_bits[1])
+        else:
+            raise TypeError("Invalid num_bits. num_bits must be a postivie integer or tuple (4,3).")
+
+        mse = ((quant_centers - centers)**2 * counts).mean()
+
+        mses.append(mse.cpu())
+        arguments.append(i)
+
+    logging.debug("mses={}".format(mses))
+    argmin = np.argmin(mses)
+    calib_amax = centers[arguments[argmin]]
+
+    return calib_amax
+
+
+def _compute_amax_percentile(calib_hist, calib_bin_edges, percentile):
+    """Returns amax that clips the percentile fraction of collected data"""
+
+    if percentile < 0 or percentile > 100:
+        raise ValueError("Invalid percentile. Must be in range 0 <= percentile <= 100.")
+
+    # If calibrator hasn't collected any data, return none
+    if calib_bin_edges is None and calib_hist is None:
+        return None
+
+    total = calib_hist.sum()
+    cdf = np.cumsum(calib_hist / total)
+    idx = np.searchsorted(cdf, percentile / 100)
+    calib_amax = calib_bin_edges[idx]
+    calib_amax = torch.tensor(calib_amax.item())  #pylint: disable=not-callable
+
+    return calib_amax
+
+
+def calibrate_weights(model, method="percentile", perchannel=True, percentile=99.99, num_bins=2048):
+    """Calibrate weights of all child quantized modules
+
+    Ideally, we would split calibration functionality to histogram collector and calibrator which
+    takes histogram and compute amax. But since we haven't decoupled collector and calibrator, it
+    is easier to create a separate function to calibrate weight.
+
+    .. note::
+        This function uses `method` specified by the argument to decide which method to use, NOT the one
+        specified by the calibrator embedded in weight_quantizer.
+        We haven't moved calibration to GPU, so everything is transfered to CPU
+
+    Args:
+        model: A torch.nn.Module.
+        method: A string of calibration method. Supports "mse" and "percentile". Default "percentile"
+        perchannel: A bool. Set channel/neuron axis if True. Default True.
+        percentile: A float. Default 99.99
+        num_bins: A integer. Number of bins of histogram. Default 2048.
+
+    """
+    for name, module in model.named_modules():
+        if hasattr(module, "weight") and hasattr(module, "weight_quantizer"):
+            logging.info("Calibrate weight of %s", name)
+            num_bits = module.weight_quantizer.num_bits
+            unsigned = module.weight_quantizer.unsigned
+            channel_second_modules = (quant_nn.QuantConvTranspose1d, quant_nn.QuantConvTranspose2d,
+                                      quant_nn.QuantConvTranspose3d)
+            if perchannel:
+                axis = 1 if isinstance(module, channel_second_modules) else 0
+            else:
+                axis = None
+            axis_size = module.weight.shape[axis] if axis is not None else 1
+
+            # Histogram is always collected even if method is "max". Although "max" is supported here
+            # but it is not the primary usage of this function
+            if axis is None:
+                input_weights = module.weight.abs().cpu().detach().numpy()
+                calib_hist, calib_bin_edges = np.histogram(input_weights, bins=2048, range=(0, input_weights.max()))
+                calib_hist = [calib_hist]
+                calib_bin_edges = [calib_bin_edges]
+            else:
+                calib_hist = []
+                calib_bin_edges = []
+                for i in range(axis_size):
+                    input_weights = module.weight.index_select(axis, torch.tensor(
+                        i, device=module.weight.device)).abs().cpu().detach().numpy()
+                    hist, bin_edges = np.histogram(input_weights, bins=num_bins, range=(0, input_weights.max()))
+                    calib_hist.append(hist)
+                    calib_bin_edges.append(bin_edges)
+
+            calib_amax = []
+            if method == "max":
+                reduce_axis = list(range(module.weight.dim()))
+                reduce_axis.remove(axis)
+                calib_amax.append(quant_utils.reduce_amax(module.weight, axis=reduce_axis))
+            elif method == 'mse':
+                for i in range(axis_size):
+                    calib_amax.append(_compute_amax_mse(calib_hist[i], calib_bin_edges[i], num_bits, unsigned))
+            elif method == 'percentile':
+                for i in range(axis_size):
+                    calib_amax.append(_compute_amax_percentile(calib_hist[i], calib_bin_edges[i], percentile))
+            else:
+                raise TypeError("Unsupported calibration method {}".format(method))
+
+            if axis is None:
+                calib_amax = calib_amax[0]
+            else:
+                calib_amax_shape = [1] * module.weight.dim()
+                calib_amax_shape[axis] = module.weight.shape[axis]
+                calib_amax = torch.stack(calib_amax).reshape(calib_amax_shape)
+            module.weight_quantizer.amax = calib_amax.detach().cpu().numpy()
+

tools/pytorch-quantization/pytorch_quantization/calib/max.py

@@ -0,0 +1,111 @@
+#
+# 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.
+#
+
+
+"""Calibrator that returns the absolute max of all collected tensors"""
+from absl import logging
+import torch
+
+from pytorch_quantization.calib.calibrator import _Calibrator
+from pytorch_quantization import utils as quant_utils
+
+class MaxCalibrator(_Calibrator):
+    """Max calibrator, tracks the maximum value globally
+
+    Args:
+        calib_desc: A MaxCalibDescriptor.
+        num_bits: An integer. Number of bits of quantization.
+        axis: A tuple. see QuantDescriptor.
+        unsigned: A boolean. using unsigned quantization.
+
+    Readonly Properties:
+        amaxs: A list of amax. Numpy array is saved as it is likely to be used for some plot.
+    """
+    def __init__(self, num_bits, axis, unsigned, track_amax=False):
+        super(MaxCalibrator, self).__init__(num_bits, axis, unsigned)
+        self._track_amax = track_amax
+        if self._track_amax:
+            self._amaxs = []  # shall we have a better name?
+        self._calib_amax = None
+
+    # pylint:disable=missing-docstring
+    @property
+    def amaxs(self):
+        return self._amaxs
+    # pylint:enable=missing-docstring
+
+    def collect(self, x):
+        """Tracks the absolute max of all tensors
+
+        Args:
+            x: A tensor
+
+        Raises:
+            RuntimeError: If amax shape changes
+        """
+        if torch.min(x) < 0.:
+            logging.log_first_n(
+                logging.INFO,
+                ("Calibrator encountered negative values. It shouldn't happen after ReLU. "
+                 "Make sure this is the right tensor to calibrate."),
+                1)
+            x = x.abs()
+
+        # Swap axis to reduce.
+        axis = self._axis if isinstance(self._axis, (list, tuple)) else [self._axis]
+        # Handle negative axis.
+        axis = [x.dim() + i if isinstance(i, int) and i < 0 else i for i in axis]
+        reduce_axis = []
+        for i in range(x.dim()):
+            if not i in axis:
+                reduce_axis.append(i)
+        local_amax = quant_utils.reduce_amax(x, axis=reduce_axis).detach()
+        if self._calib_amax is None:
+            self._calib_amax = local_amax
+        else:
+            if local_amax.shape != self._calib_amax.shape:
+                raise RuntimeError("amax shape changed!")
+            self._calib_amax.copy_(torch.max(self._calib_amax, local_amax).data)
+
+        if self._track_amax:
+            self._amaxs.append(local_amax.cpu().numpy())
+
+    def reset(self):
+        """Reset the collected absolute max"""
+        self._calib_amax = None
+
+    def compute_amax(self):
+        """Return the absolute max of all tensors collected"""
+        return self._calib_amax
+
+    # pylint:disable=missing-docstring
+    def __str__(self):
+        s = "MaxCalibrator("
+        s += "track_amax={_track_amax}"
+        s += ")"
+        return s.format(**self.__dict__)
+
+    def __repr__(self):
+        s = "MaxCalibrator("
+        s += super(MaxCalibrator, self).__repr__()
+        s += " calib_amax={_calib_amax}"
+        s += " track_amax={_track_amax}"
+        if self._track_amax:
+            s += " amaxs={_amaxs}"
+        s += ")"
+        return s.format(**self.__dict__)
+    # pylint:enable=missing-docstring

tools/pytorch-quantization/pytorch_quantization/nn/__init__.py

@@ -0,0 +1,25 @@
+#
+# 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.
+#
+
+
+from pytorch_quantization.nn.modules.tensor_quantizer import *
+from pytorch_quantization.nn.modules.quant_conv import *
+from pytorch_quantization.nn.modules.quant_linear import *
+from pytorch_quantization.nn.modules.quant_pooling import *
+from pytorch_quantization.nn.modules.clip import *
+from pytorch_quantization.nn.modules.quant_rnn import *
+from pytorch_quantization.nn.modules.quant_instancenorm import *

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

@@ -0,0 +1,286 @@
+#
+# 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/_functions/rnn.py
+with following modification
+    fusedBackend is removed
+    CudnnRNN is removed
+    Hack for ONNX in RNN() is removed
+Only LSTM is quantized. Other paths are excluded in __all__
+"""
+
+import warnings
+from torch.autograd import NestedIOFunction
+from torch.nn import functional as F
+import torch
+import itertools
+from functools import partial
+
+__all__ = ["LSTMCell", "RNN"]
+
+def RNNReLUCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None):
+    hy = F.relu(F.linear(input, w_ih, b_ih) + F.linear(hidden, w_hh, b_hh))
+    return hy
+
+
+def RNNTanhCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None):
+    hy = torch.tanh(F.linear(input, w_ih, b_ih) + F.linear(hidden, w_hh, b_hh))
+    return hy
+
+
+def LSTMCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, input_quantizer=None, weight_quantizer=None):
+    """Quantized LSTM Cell
+
+    The assumption is at inference time, only one fused gemm will be launched for one time step Weights of 4 gates
+    are fused together, and activation from layer and recurrent paths are fused togather. ``input_quantizer`` will be
+    applied on the fused activation tensor. And ``weight_quantizer`` will be applied on the fused weight tensor.
+    """
+
+    hx, cx = hidden
+    if input_quantizer is not None:
+        input, hx = input_quantizer(torch.cat([input, hx], 1)).split([input.size()[1], hx.size()[1]], 1)
+    if weight_quantizer is not None:
+        w_ih, w_hh = weight_quantizer(torch.cat([w_ih, w_hh], 1)).split([w_ih.size()[1], w_hh.size()[1]], 1)
+    gates = F.linear(input, w_ih, b_ih) + F.linear(hx, w_hh, b_hh)
+
+    ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)
+
+    ingate = torch.sigmoid(ingate)
+    forgetgate = torch.sigmoid(forgetgate)
+    cellgate = torch.tanh(cellgate)
+    outgate = torch.sigmoid(outgate)
+
+    cy = (forgetgate * cx) + (ingate * cellgate)
+    hy = outgate * torch.tanh(cy)
+
+    return hy, cy
+
+
+def GRUCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None):
+    gi = F.linear(input, w_ih, b_ih)
+    gh = F.linear(hidden, w_hh, b_hh)
+    i_r, i_i, i_n = gi.chunk(3, 1)
+    h_r, h_i, h_n = gh.chunk(3, 1)
+
+    resetgate = torch.sigmoid(i_r + h_r)
+    inputgate = torch.sigmoid(i_i + h_i)
+    newgate = torch.tanh(i_n + resetgate * h_n)
+    hy = newgate + inputgate * (hidden - newgate)
+
+    return hy
+
+
+def StackedRNN(inners, num_layers, lstm=False, dropout=0, train=True):
+
+    num_directions = len(inners)
+    total_layers = num_layers * num_directions
+
+    def forward(input, hidden, weight, batch_sizes, input_quantizers, weight_quantizers):
+        assert(len(weight) == total_layers)
+        next_hidden = []
+
+        if lstm:
+            hidden = list(zip(*hidden))
+
+        for i in range(num_layers):
+            all_output = []
+            for j, inner in enumerate(inners):
+                l = i * num_directions + j
+
+                hy, output = inner(input, hidden[l], weight[l], batch_sizes,
+                                   input_quantizer=input_quantizers[l], weight_quantizer=weight_quantizers[l])
+                next_hidden.append(hy)
+                all_output.append(output)
+
+            input = torch.cat(all_output, input.dim() - 1)
+
+            if dropout != 0 and i < num_layers - 1:
+                input = F.dropout(input, p=dropout, training=train, inplace=False)
+
+        if lstm:
+            next_h, next_c = zip(*next_hidden)
+            next_hidden = (
+                torch.cat(next_h, 0).view(total_layers, *next_h[0].size()),
+                torch.cat(next_c, 0).view(total_layers, *next_c[0].size())
+            )
+        else:
+            next_hidden = torch.cat(next_hidden, 0).view(
+                total_layers, *next_hidden[0].size())
+
+        return next_hidden, input
+
+    return forward
+
+
+def Recurrent(inner, reverse=False):
+    def forward(input, hidden, weight, batch_sizes, input_quantizer, weight_quantizer):
+        output = []
+        steps = range(input.size(0) - 1, -1, -1) if reverse else range(input.size(0))
+        for i in steps:
+            hidden = inner(input[i], hidden, *weight,
+                           input_quantizer=input_quantizer, weight_quantizer=weight_quantizer)
+            # hack to handle LSTM
+            output.append(hidden[0] if isinstance(hidden, tuple) else hidden)
+
+        if reverse:
+            output.reverse()
+        output = torch.cat(output, 0).view(input.size(0), *output[0].size())
+
+        return hidden, output
+
+    return forward
+
+
+def variable_recurrent_factory(inner, reverse=False):
+    if reverse:
+        return VariableRecurrentReverse(inner)
+    else:
+        return VariableRecurrent(inner)
+
+
+def VariableRecurrent(inner):
+    def forward(input, hidden, weight, batch_sizes, input_quantizer, weight_quantizer):
+
+        output = []
+        input_offset = 0
+        last_batch_size = batch_sizes[0]
+        hiddens = []
+        flat_hidden = not isinstance(hidden, tuple)
+        if flat_hidden:
+            hidden = (hidden,)
+        for batch_size in batch_sizes:
+            step_input = input[input_offset:input_offset + batch_size]
+            input_offset += batch_size
+
+            dec = last_batch_size - batch_size
+            if dec > 0:
+                hiddens.append(tuple(h[-dec:] for h in hidden))
+                hidden = tuple(h[:-dec] for h in hidden)
+            last_batch_size = batch_size
+
+            if flat_hidden:
+                hidden = (inner(step_input, hidden[0], *weight,
+                                input_quantizer=input_quantizer, weight_quantizer=weight_quantizer),)
+            else:
+                hidden = inner(step_input, hidden, *weight,
+                               input_quantizer=input_quantizer, weight_quantizer=weight_quantizer)
+
+            output.append(hidden[0])
+        hiddens.append(hidden)
+        hiddens.reverse()
+
+        hidden = tuple(torch.cat(h, 0) for h in zip(*hiddens))
+        assert hidden[0].size(0) == batch_sizes[0]
+        if flat_hidden:
+            hidden = hidden[0]
+        output = torch.cat(output, 0)
+
+        return hidden, output
+
+    return forward
+
+
+def VariableRecurrentReverse(inner):
+    def forward(input, hidden, weight, batch_sizes, input_quantizer, weight_quantizer):
+        output = []
+        input_offset = input.size(0)
+        last_batch_size = batch_sizes[-1]
+        initial_hidden = hidden
+        flat_hidden = not isinstance(hidden, tuple)
+        if flat_hidden:
+            hidden = (hidden,)
+            initial_hidden = (initial_hidden,)
+        hidden = tuple(h[:batch_sizes[-1]] for h in hidden)
+        for i in reversed(range(len(batch_sizes))):
+            batch_size = batch_sizes[i]
+            inc = batch_size - last_batch_size
+            if inc > 0:
+                hidden = tuple(torch.cat((h, ih[last_batch_size:batch_size]), 0)
+                               for h, ih in zip(hidden, initial_hidden))
+            last_batch_size = batch_size
+            step_input = input[input_offset - batch_size:input_offset]
+            input_offset -= batch_size
+
+            if flat_hidden:
+                hidden = (inner(step_input, hidden[0], *weight,
+                                input_quantizer=input_quantizer, weight_quantizer=weight_quantizer),)
+            else:
+                hidden = inner(step_input, hidden, *weight,
+                               input_quantizer=input_quantizer, weight_quantizer=weight_quantizer)
+            output.append(hidden[0])
+
+        output.reverse()
+        output = torch.cat(output, 0)
+        if flat_hidden:
+            hidden = hidden[0]
+        return hidden, output
+
+    return forward
+
+
+def AutogradRNN(mode, input_size, hidden_size, num_layers=1, batch_first=False,
+                dropout=0, train=True, bidirectional=False, variable_length=False,
+                dropout_state=None, flat_weight=None,
+                input_quantizers=None, weight_quantizers=None):
+
+    if mode == 'RNN_RELU':
+        cell = RNNReLUCell
+    elif mode == 'RNN_TANH':
+        cell = RNNTanhCell
+    elif mode == 'LSTM':
+        cell = LSTMCell
+    elif mode == 'GRU':
+        cell = GRUCell
+    else:
+        raise Exception('Unknown mode: {}'.format(mode))
+
+    rec_factory = variable_recurrent_factory if variable_length else Recurrent
+
+    if bidirectional:
+        layer = (rec_factory(cell), rec_factory(cell, reverse=True))
+    else:
+        layer = (rec_factory(cell),)
+
+    func = StackedRNN(layer,
+                      num_layers,
+                      (mode == 'LSTM'),
+                      dropout=dropout,
+                      train=train)
+
+    def forward(input, weight, hidden, batch_sizes, input_quantizers, weight_quantizers):
+        if batch_first and not variable_length:
+            input = input.transpose(0, 1)
+
+        nexth, output = func(input, hidden, weight, batch_sizes, input_quantizers, weight_quantizers)
+
+        if batch_first and not variable_length:
+            output = output.transpose(0, 1)
+
+        return output, nexth
+
+    return forward
+
+
+def RNN(*args, **kwargs):
+
+    def forward(input, *fargs, **fkwargs):
+        func = AutogradRNN(*args, **kwargs)
+        return func(input, *fargs, **fkwargs)
+
+    return forward

tools/pytorch-quantization/pytorch_quantization/nn/functional.py

@@ -0,0 +1,61 @@
+#
+# 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 supportive functions"""
+from absl import logging
+
+import torch
+from torch.autograd import Function
+
+
+class ClipFunction(Function):
+    """An universal tensor clip function
+
+    Pytorch's clamp() only supports scalar range and doesn't support broadcast. This implementation uses min/max which
+    is more genaral. The gradient is defined according to IBM's PACT paper https://arxiv.org/abs/1805.06085, which is
+    also the behavior of Tensorflow's clip_by_value()
+    """
+
+    @staticmethod
+    def forward(ctx, input, clip_value_min, clip_value_max):
+        output = torch.min(input, clip_value_max)
+        output = torch.max(output, clip_value_min)
+        ctx.save_for_backward(input, clip_value_min, clip_value_max)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, clip_value_min, clip_value_max = ctx.saved_tensors
+        min_mask = (input > clip_value_min).to(grad_output.dtype)
+        max_mask = (input < clip_value_max).to(grad_output.dtype)
+        grad_input = grad_output * min_mask * max_mask
+
+        if clip_value_min.requires_grad or clip_value_max.requires_grad:
+            logging.log_first_n(logging.WARNING, "Learning clip min/max is experimental, use at your own risk :).", 1)
+        if clip_value_min.numel() != 1 or clip_value_max.numel() != 1:
+            raise ValueError("Learnable min/max can only be scalar, got size %s and %s." % (clip_value_min.size(),
+                                                                                            clip_value_max.size()))
+
+        # Ensure the dtypes of min/max grads matches the input dtype
+        # This might be necessary if running w/ AMP which will cast to fp32 before `sum()`
+        grad_clip_value_min = (grad_output * (1. - min_mask)).sum().to(clip_value_min.dtype) if clip_value_min.requires_grad else None
+        grad_clip_value_max = (grad_output * (1. - max_mask)).sum().to(clip_value_min.dtype) if clip_value_max.requires_grad else None
+
+        return grad_input, grad_clip_value_min, grad_clip_value_max
+
+
+clip = ClipFunction.apply

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)

tools/pytorch-quantization/pytorch_quantization/optim/helper.py

@@ -0,0 +1,131 @@
+#
+# 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.
+#
+
+
+"""Helper functions for quant optimizer/trainer"""
+
+import re
+
+from absl import logging
+
+def match_parameters(model, patterns):
+    """Returns an generator over module parameters if name matches key
+
+    It is useful to group parameters, and apply different functions to different group. This function provides an easy
+    way to group them.
+
+    Args:
+        model: A Module
+        patterns: A list of strings that will be used to match parameter names. If parameter name contains any pattern,
+            it will be yield
+
+    Yields:
+        param: Module parameters
+    """
+    for name, param in model.named_parameters():
+        for pattern in patterns:
+            if re.search(pattern, name):
+                yield param
+
+def group_parameters(model, patterns_list, lrs=None, momentums=None, weight_decays=None):
+    """Group parameters for using per-parameters option in optimizer
+
+    Returns a list of dict that matches Pytorch optimizer fashion, see
+    https://pytorch.org/docs/stable/optim.html#per-parameter-options for more details.
+
+    Example:
+        >>> [
+        >>>    {'params': model.base.parameters()},
+        >>>    {'params': model.classifier.parameters(), 'lr': 1e-3}
+        >>> ]
+    Parameters will be grouped w.r.t first level of the keys_list. e.g. `keys_list=[['conv1', 'conv2'], ['conv3']]` will
+    return 2 groups, one with `conv1` and `conv2` in name, and the other with `conv3` in name.
+
+    If lr, momentum or weight_decay are supplied, they will be added to the group as well.
+
+
+    Args:
+        model: A module
+        patterns_list: A list of list of strings. WARNING: patters must be EXCLUSIVE, the function doesn't
+            perform exclusive check.
+        lrs: A list of float with same length as keys_list or None.
+        momentums: A list of float with same length as keys_list or None.
+        weight_decays: A list of float with same length as keys_list or None.
+
+    Returns:
+        param_group: A list of dict
+
+    """
+    param_groups = []
+    for pattern in patterns_list:
+        if not isinstance(pattern, list):
+            raise TypeError("patterns_list must be list of list of patterns")
+        param_groups.append({'params': match_parameters(model, pattern)})
+
+    if lrs is not None:
+        if len(lrs) != len(patterns_list):
+            raise TypeError("len(lrs) must match len(patterns_list)")
+        for i, lr in enumerate(lrs):
+            param_groups[i]['lr'] = lr
+
+    if momentums is not None:
+        if len(momentums) != len(patterns_list):
+            raise TypeError("len(momentums) must match len(patterns_list)")
+        for i, momentum in enumerate(momentums):
+            param_groups[i]['momentum'] = momentum
+
+    if weight_decays is not None:
+        if len(weight_decays) != len(patterns_list):
+            raise TypeError("len(weight_decays) must match len(patterns_list)")
+        for i, weight_decay in enumerate(weight_decays):
+            param_groups[i]['weight_decay'] = weight_decay
+
+    return param_groups
+
+def freeze_parameters(model, patterns):
+    """Set requires_grad to False if patterns match name
+
+    Args:
+        model: A Module
+        patterns: A list of strings that will be used to match parameter names. If parameter name contains any pattern,
+            it will be frozen.
+    """
+    for name, param in model.named_parameters():
+        for pattern in patterns:
+            if re.search(pattern, name):
+                logging.warning("Freeze %s.", name)
+                param.requires_grad = False
+
+def quant_weight_inplace(model):
+    """Make quantization inplace
+
+    Search for quantized modules including QuantConvNd and QuantLinear, make weight quantization in place using
+    weight_quantizer.
+
+    Most publications of quantization aware training uses STE by default, which is really an approximation of
+    derivative of the nondifferentiable quantization function, which works to some extended but by no means the F=ma of
+    the problem.
+    Inplace quantization can be used to implement relax-and-round, which is a common method in Discrete Optimization's
+    or Integer Programming.
+    """
+    for name, module in model.named_modules():
+        if hasattr(module, '_weight_quantizer') and module.weight_quantizer is not None:
+            if not module.weight_quantizer.fake_quant:
+                logging.warning(("In-place real quantization is VERY dangerous and should be used for inference only. "
+                                 "Make sure that is the desired behavior."))
+            logging.warning("In-place quantize weight of %s", name)
+            module.weight.data.copy_(module.weight_quantizer(module.weight))

tools/pytorch-quantization/pytorch_quantization/quant_modules.py

@@ -0,0 +1,182 @@
+#
+# 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.
+#
+"""Dynamically replace the modules with quantized versions."""
+
+from collections import namedtuple
+from contextlib import contextmanager
+
+import torch
+from pytorch_quantization import nn as quant_nn
+
+__all__ = ['initialize', 'deactivate', 'enable_onnx_export']
+
+# Definition of the named tuple that is used to store mapping of the quantized modules
+_quant_entry = namedtuple('quant_entry', 'orig_mod mod_name replace_mod')
+
+# Global member of the file that contains the mapping of quantized modules
+_DEFAULT_QUANT_MAP = [
+    _quant_entry(torch.nn, "Conv1d", quant_nn.QuantConv1d),
+    _quant_entry(torch.nn, "Conv2d", quant_nn.QuantConv2d),
+    _quant_entry(torch.nn, "Conv3d", quant_nn.QuantConv3d),
+    _quant_entry(torch.nn, "ConvTranspose1d", quant_nn.QuantConvTranspose1d),
+    _quant_entry(torch.nn, "ConvTranspose2d", quant_nn.QuantConvTranspose2d),
+    _quant_entry(torch.nn, "ConvTranspose3d", quant_nn.QuantConvTranspose3d),
+    _quant_entry(torch.nn, "Linear", quant_nn.QuantLinear),
+    _quant_entry(torch.nn, "LSTM", quant_nn.QuantLSTM),
+    _quant_entry(torch.nn, "LSTMCell", quant_nn.QuantLSTMCell),
+    _quant_entry(torch.nn, "AvgPool1d", quant_nn.QuantAvgPool1d),
+    _quant_entry(torch.nn, "AvgPool2d", quant_nn.QuantAvgPool2d),
+    _quant_entry(torch.nn, "AvgPool3d", quant_nn.QuantAvgPool3d),
+    _quant_entry(torch.nn, "AdaptiveAvgPool1d", quant_nn.QuantAdaptiveAvgPool1d),
+    _quant_entry(torch.nn, "AdaptiveAvgPool2d", quant_nn.QuantAdaptiveAvgPool2d),
+    _quant_entry(torch.nn, "AdaptiveAvgPool3d", quant_nn.QuantAdaptiveAvgPool3d),
+]
+
+
+class QuantModuleReplacementHelper():
+    """To help replace torch.nn modules with quantized versions.
+
+    This module is used to replace (by monkey patching) the torch.nn modules with their
+    quantized versions as provided by either tool's internal implementation or any other
+    user provided custom module.
+
+    Attributes:
+        orginal_func_map: A dict. Maintains the original torch.nn module mapping.
+        quant_support_list: A list. Contains the names of modules for which a quantized
+            version is provided by the tool.
+        quant_map: A dict. Contains the map of the module name and its quantized versions.
+        quant_switch_opt: A dict. A map to indicate which modules to be left unreplaced with
+            their quantized versions. This dict is updated by a list provided from the user
+            which indicates the modules to leave out in monkey patching.
+
+    """
+
+    def __init__(self):
+
+        # Will hold the original modules to be replaced back
+        self.orginal_func_map = set()
+
+        # Maintains the list of supported quantized modules by the tool as default
+        self.default_quant_map = _DEFAULT_QUANT_MAP
+
+        # Will hold the final quantized modules after checking if user supplied any
+        # custom quantized functions.
+        self.quant_map = set()
+
+    def prepare_state(self, float_module_list=None, custom_map=None):
+        """
+        Prepare the internal variables that would used in the monkey patching mechanism later.
+        1. Set up the list of quantized modules that are supported by the tool for torch.nn.
+        2. Set up the custom mapping for modules other than torch.nn.
+        3. Use the float_module_list to switch off the monkey patching replacement for user indicated modules
+        """
+
+        # For the default quantized modules supported, generate the quant_map
+        for item in self.default_quant_map:
+            if float_module_list is not None and item.mod_name in float_module_list:
+                # Skip this module if this is present in the float_module_list
+                continue
+            else:
+                # append the modules into the variable that will be used in monkey patching
+                self.quant_map.add(item)
+                # also store the original module to be used in reverse monkey patching
+                self.orginal_func_map.add(
+                    _quant_entry(item.orig_mod, item.mod_name, getattr(item.orig_mod, item.mod_name)))
+
+        # Add custom modules to the quant_map
+        if custom_map is not None:
+            for item in custom_map:
+                # append the custom modules to the list that will be used in monkey patching
+                # Note that we convert a tuple to a named tuple here
+                self.quant_map.add(_quant_entry(item[0], item[1], item[2]))
+                # also store the original module in another list which will be used to reverse monkey patching
+                self.orginal_func_map.add(_quant_entry(item[0], item[1], getattr(item[0], item[1])))
+
+    def apply_quant_modules(self):
+        """
+        For the modules registered in the quant_map, simply monkey patch them and also store the
+        original modules so that they could be later replaced back.
+        """
+        for entry in self.quant_map:
+            setattr(entry.orig_mod, entry.mod_name, entry.replace_mod)
+
+    def restore_float_modules(self):
+        """
+        Reverse the effect of monkey patch by using the orginal_func_map to replace back the
+        original modules.
+        """
+        for entry in self.orginal_func_map:
+            setattr(entry.orig_mod, entry.mod_name, entry.replace_mod)
+
+
+def initialize(float_module_list=None, custom_quant_modules=None):
+    """Dynamic module replacement using monkey patching.
+
+    Dynamically monkey patches the modules with their quantized versions. Internally, the
+    state is maintained by a helper class object which helps in replacing the original
+    modules back.
+
+    Args:
+        float_module_list: A list. User supplied list which indicates which modules to not monkey patch.
+        custom_quant_modules: A dict. A mapping provided by user to indicate any other module apart
+            from torch.nn and its corresponding quantized version.
+
+    Returns:
+        nothing.
+
+    Typical usage example:
+
+        # Define the deny list for torch.nn modules and custom map for modules other than torch.nn.
+        float_module_list = ["Linear"]
+        custom_quant_modules = [(torch.nn, "Linear", quant_nn.QuantLinear)]
+        ## Monkey patch the modules
+        pytorch_quantization.quant_modules.initialize(float_module_list, custom_modules)
+        ## Use the quantized modules
+        pytorch_quantization.quant_modules.deactivate()
+    """
+    _quant_module_helper_object.prepare_state(float_module_list, custom_quant_modules)
+    _quant_module_helper_object.apply_quant_modules()
+
+
+def deactivate():
+    """Dynamic module replacement which reverses the monkey patching.
+
+    Dynamically replaces back the original modules that were monkey patched earlier
+    in the initialize() function call using helper class object which maintains the state.
+    """
+    _quant_module_helper_object.restore_float_modules()
+
+
+# Global object that maintains the state of the modules that are replaced.
+_quant_module_helper_object = QuantModuleReplacementHelper()
+
+
+@contextmanager
+def enable_onnx_export():
+    """Context manager to enable onnx export.
+    
+    .. code-block:: python
+
+        with pytorch_quantization.enable_onnx_export():
+            # export onnx model
+            torch.onnx.export(model, ...)
+    
+    """
+    quant_nn.TensorQuantizer._enable_onnx_export = True
+    yield
+
+    quant_nn.TensorQuantizer._enable_onnx_export = False

tools/pytorch-quantization/pytorch_quantization/tensor_quant.py

@@ -0,0 +1,624 @@
+#
+# 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.
+#
+"""Basic tensor quantization functions"""
+import numpy as np
+import numbers
+import yaml
+
+from absl import logging
+
+import torch
+import torch._C._onnx as _C_onnx
+from torch.autograd import Function
+
+from pytorch_quantization import cuda_ext
+
+from torch.onnx import symbolic_helper
+
+
+class ScaledQuantDescriptor():
+    """Supportive descriptor of quantization
+
+    Describe how a tensor should be quantized. A QuantDescriptor and a tensor defines a quantized tensor.
+
+    Args:
+        num_bits: An integer or a tuple of two integers. 
+            Specifically, `num_bits` can be:
+            
+            #. A positive integer argument for integer qunatization. `num_bits` specify 
+                the number of bits used for integer quantization.
+
+            #. Constant integer tuple (4,3) for E4M3 floating point quantization emulating
+                Nvidia's FP8 quantization. E4M3 quantization only supports per-tensor quantization.
+
+            Default: 8.
+        name: Seems a nice thing to have
+
+    Keyword Arguments:
+        fake_quant: A boolean. If True, use fake quantization mode. Default True.
+        axis: None, int or tuple of int. axes which will have its own max for computing scaling factor.
+            If None (the default), use per tensor scale.
+            Must be in the range [-rank(input_tensor), rank(input_tensor)).
+            e.g. For a KCRS weight tensor, quant_axis=(0) will yield per channel scaling.
+            Default None.
+        amax: A float or list/ndarray of floats of user specified absolute max range. If supplied,
+            ignore quant_axis and use this to quantize. If learn_amax is True, will be used to initialize
+            learnable amax. Default None.
+        learn_amax: A boolean. If True, learn amax. Default False.
+        scale_amax: A float. If supplied, multiply amax by scale_amax. Default None. It is useful for some
+            quick experiment.
+        calib_method: A string. One of ["max", "histogram"] indicates which calibration to use. Except the simple
+            max calibration, other methods are all hisogram based. Default "max".
+        unsigned: A Boolean. If True, use unsigned. Default False.
+
+    Raises:
+        TypeError: If unsupported type is passed in.
+
+    Read-only properties:
+        - fake_quant:
+        - name:
+        - learn_amax:
+        - scale_amax:
+        - axis:
+        - calib_method:
+        - num_bits:
+        - amax:
+        - unsigned:
+    """
+
+    def __init__(self, num_bits=8, name=None, **kwargs):
+        if isinstance(num_bits, int):
+            if num_bits < 0:
+                raise ValueError("num_bits must be > 0, not {}.".format(num_bits))
+            if num_bits == 0:
+                logging.error("num_bits is 0. This will result in the tensor being quantized to all zeros."
+                              " This mode should only be used for debugging purposes.")
+        elif num_bits != (4, 3):
+            raise TypeError("num_bits must be a postive integer or tuple (4,3), not {}.".format(type(num_bits)))
+
+        self._num_bits = num_bits
+        if not isinstance(name, str) and name is not None:
+            raise TypeError("name must be a string or None, not {}.".format(type(name)))
+        self._name = name
+
+        self._fake_quant = kwargs.pop('fake_quant', True)
+        self._axis = kwargs.pop('axis', None)
+        if self._axis is not None:
+            logging.debug("Meaning of axis has changed since v2.0. Make sure to update.")
+        self._learn_amax = kwargs.pop('learn_amax', False)
+        if self._learn_amax and self._axis is not None:
+            raise TypeError("axis is ignored and must be None when learn_amax is true, got {}.".format(type(
+                self._axis)))
+        amax = kwargs.pop('amax', None)
+        if amax is not None:
+            if not isinstance(amax, float) and not isinstance(amax, list) and not isinstance(amax, np.ndarray):
+                raise TypeError("amax must be float, list or ndarray, not {}".format(type(amax)))
+            # Make it single precision array
+            self._amax = np.array(amax, dtype=np.float32)
+        else:
+            self._amax = amax
+
+        self._scale_amax = kwargs.pop('scale_amax', None)
+        self._calib_method = kwargs.pop('calib_method', "max")
+        self._unsigned = kwargs.pop('unsigned', False)
+        self._narrow_range = kwargs.pop('narrow_range', False)
+
+        if kwargs:
+            raise TypeError("Unused keys: {}".format(kwargs.keys()))
+
+    # pylint:disable=missing-docstring
+    @property
+    def num_bits(self):
+        return self._num_bits
+
+    @property
+    def fake_quant(self):
+        return self._fake_quant
+
+    @property
+    def axis(self):
+        return self._axis
+
+    @property
+    def amax(self):
+        return self._amax
+
+    @property
+    def learn_amax(self):
+        return self._learn_amax
+
+    @property
+    def scale_amax(self):
+        return self._scale_amax
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def calib_method(self):
+        return self._calib_method
+
+    @property
+    def unsigned(self):
+        return self._unsigned
+
+    @property
+    def narrow_range(self):
+        return self._narrow_range
+
+    # pylint:enable=missing-docstring
+
+    def __str__(self):
+        s = (self._name + ': ') if self._name is not None else 'QuantDescriptor'
+        s += "({}{}bit".format("unsigned " if self._unsigned else "", self._num_bits)
+        s += " fake" if self._fake_quant else " real"
+        s += " axis={}".format(self._axis if self._axis is not None else " per-tensor")
+        if isinstance(self._amax, torch.Tensor):
+            s += " amax={}".format(
+                np.array2string(self._amax.cpu().numpy().flatten(), edgeitems=1, formatter={'all': "{:.2e}".format}))
+        elif self._amax is not None:
+            s += " amax={_amax}"
+            s += " full_range"
+        if self._learn_amax:
+            s += " learn_amax"
+        if self._scale_amax:
+            s += " scale_amax={_scale_amax}"
+        s += ")"
+        return s.format(**self.__dict__)
+
+    def __eq__(self, rhs):
+        """Compare 2 descriptors"""
+        return self.__dict__ == rhs.__dict__
+
+    def dict(self):
+        """Serialize to dict
+
+        The build-in __dict__ method returns all the attributes, which includes those have default value and have
+        protected prefix "_". This method only returns those have values other than the default one and don't have _ in
+        key. Construct a instance by dict returned by this method should get exactly the same instance.
+        """
+        obj_dict = {}
+        obj_dict['num_bits'] = self._num_bits
+        obj_dict['name'] = self._name
+
+        if not self._fake_quant:
+            obj_dict['fake_quant'] = self._fake_quant
+        if self._axis is not None:
+            obj_dict['axis'] = self._axis
+        if self._amax is not None:
+            obj_dict['amax'] = self._amax.tolist()
+        if self._scale_amax is not None:
+            obj_dict['scale_amax'] = self._scale_amax
+        if self._learn_amax:
+            obj_dict['learn_amax'] = self._learn_amax
+        if self._unsigned:
+            obj_dict['unsigned'] = self._unsigned
+
+        return obj_dict
+
+    def to_yaml(self):
+        """Create yaml serialization
+        Some attributes need special treatment to have human readable form, including amax, axis.
+        """
+        obj_dict = self.dict()
+
+        if "axis" in obj_dict:
+            obj_dict['axis'] = list(obj_dict['axis'])
+        return yaml.dump(obj_dict, width=120)
+
+    @classmethod
+    def from_yaml(cls, yaml_str):
+        """Create descriptor from yaml str"""
+        obj_dict = yaml.safe_load(yaml_str)
+
+        if 'axis' in obj_dict:
+            obj_dict['axis'] = tuple(obj_dict['axis'])
+
+        quant_desc = cls(**obj_dict)
+
+        return quant_desc
+
+
+QuantDescriptor = ScaledQuantDescriptor
+
+# Predefined descriptors
+QUANT_DESC_8BIT_PER_TENSOR = QuantDescriptor(num_bits=8)
+QUANT_DESC_UNSIGNED_8BIT_PER_TENSOR = QuantDescriptor(num_bits=8, unsigned=True)
+QUANT_DESC_8BIT_CONV1D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1))
+QUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1))
+QUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1))
+
+
+@torch.jit.script
+def _fake_tensor_quant_backward(inputs, amax, grad_outputs):
+    zero = grad_outputs.new_zeros(1)
+    grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero)
+    return grad_inputs
+
+
+def _onnx_int8_helper(g, inputs, amax, num_bits, unsigned, narrow_range):
+    assert num_bits == 8, "Only INT8 ONNX export is supported for now."
+    maxbound = (1 << (num_bits - 1 + int(unsigned))) - 1
+
+    if amax.numel() == 1:
+        zero_point, axis = torch.tensor(0.0, device=amax.device), None
+    else:
+        amax_init_shape = amax.shape
+        amax = amax.squeeze().data
+        assert len(amax.shape) == 1, "ONNX does not support multi-axis quantization."
+        zero_point = torch.zeros_like(amax, dtype=torch.int32).data
+        axis = list(amax_init_shape).index(list(amax.shape)[0])
+
+    zero_point = g.op("Constant", value_t=zero_point)
+
+    if not unsigned:
+        assert not narrow_range, "ONNX does not support unsigned narrow range INT8."
+        zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.INT8)
+    else:
+        zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.UINT8)
+
+    scale = amax / maxbound
+    scale.masked_fill_(scale == 0, 1.0)
+    scale = g.op("Constant", value_t=scale)
+
+    input_type = inputs.type().scalarType()
+
+    # Q inputs are currently constrained to FP32 due to a similar limitation in ORT
+    # custom ops, so cast the input if needed.
+    if input_type == "Half" or input_type == "BFloat16":
+        inputs = g.op("Cast", inputs, to_i=_C_onnx.TensorProtoDataType.FLOAT)
+
+    quantized = g.op("QuantizeLinear", inputs, scale, zero_point, axis_i=axis)
+    out = g.op("DequantizeLinear", quantized, scale, zero_point, axis_i=axis)
+
+    # DQ outputs are currently constrained to FP32 due to a similar limitation in ORT
+    # custom ops, so cast the output if needed.
+    if input_type == "Half":
+        out = g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.FLOAT16)
+    elif input_type == "BFloat16":
+        out = g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.BFLOAT16)
+
+    return out
+
+
+class FakeTensorQuantFunction(Function):
+    """Fake version of TensorQuantFunction use CUDA extension"""
+
+    @staticmethod
+    @symbolic_helper.parse_args("v", "t", "i", "b", "b")
+    def symbolic(g, inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+        return _onnx_int8_helper(g, inputs, amax, num_bits, unsigned, narrow_range)
+
+    @staticmethod
+    def forward(ctx, inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+        ctx.save_for_backward(inputs, amax)
+
+        def legacy_quant_func():
+            # The LegacyFakeTensorQuantFunction support cpu and amax with any shape that can be broadcasted to inputs.
+            outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+            return outputs / scale.to(inputs.dtype)
+
+        if not inputs.is_cuda:
+            outputs = legacy_quant_func()
+        else:
+            try:
+                if amax.numel() == 1:
+                    outputs = cuda_ext.fake_tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+                else:
+                    axis = amax.shape.index(amax.numel())
+
+                    outputs = cuda_ext.fake_tensor_quant_with_axis(inputs, amax.squeeze(), axis, num_bits, unsigned,
+                                                                   narrow_range)
+            except (RuntimeError, ValueError) as error:
+                outputs = legacy_quant_func()
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        inputs, amax = ctx.saved_tensors
+        return _fake_tensor_quant_backward(inputs, amax, grad_outputs), None, None, None, None
+
+
+def _onnx_fp8_quantize(g, inputs, scale_inv):
+    """Helper Function for Quantization"""
+    output_shape = torch.onnx.symbolic_helper._get_tensor_sizes(inputs)
+
+    # Q inputs are currently constrained to FP32 due to a similar limitation in ORT
+    # custom ops, so cast the input if needed.
+    if inputs.type().scalarType() == "Half" or inputs.type().scalarType() == "BFloat16":
+        inputs = g.op("Cast", inputs, to_i=_C_onnx.TensorProtoDataType.FLOAT)
+
+    scale = g.op("Constant", value_t=torch.tensor(scale_inv))
+    q_op = g.op("trt::TRT_FP8QuantizeLinear", inputs,
+                scale).setType(inputs.type().with_dtype(torch.uint8).with_sizes(output_shape))
+    return q_op
+
+
+def _onnx_fp8_dequantize(g, inputs, scale_inv, otype=None):
+    """Helper Function for Dequantization"""
+    output_shape = torch.onnx.symbolic_helper._get_tensor_sizes(inputs)
+
+    scale = g.op("Constant", value_t=torch.tensor(scale_inv))
+    out = g.op("trt::TRT_FP8DequantizeLinear", inputs,
+               scale).setType(inputs.type().with_dtype(torch.float32).with_sizes(output_shape))
+
+    # DQ outputs are currently constrained to FP32 due to a similar limitation in ORT
+    # custom ops, so cast the output if needed.
+    if otype == "Half":
+        out = g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.FLOAT16)
+    elif otype == "BFloat16":
+        out = g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.BFLOAT16)
+    return out
+
+
+class ScaledE4M3Function(Function):
+    """E4M3fy input with scale"""
+
+    @staticmethod
+    @symbolic_helper.parse_args("v", "t", "i", "b", "b")
+    def symbolic(g, inputs, amax=None, E=4, M=3):
+        if amax is None:
+            scale = 1.0
+        else:
+            scale = 448.0 / amax
+            scale = float(scale.masked_fill_(scale == 0, 1.0))
+
+        input_type = inputs.type().scalarType()
+        q_tensor = _onnx_fp8_quantize(g, inputs, 1.0 / scale)
+        return _onnx_fp8_dequantize(g, q_tensor, 1.0 / scale, input_type)
+
+    @staticmethod
+    def forward(ctx, inputs, amax=None, E=4, M=3):
+        if E != 4 or M != 3:
+            raise NotImplementedError("Only support E=4 & M=3 for now.")
+
+        ctx.save_for_backward(inputs)
+        ctx.amax = amax
+        zero_mask = (inputs.abs() < 1. / (1 << 24))
+
+        if amax is None:
+            outputs = cuda_ext.fake_e4m3fy(inputs)
+        else:
+            # FP8 ONNX export requires scalar `scale`.
+            # To simplify implementation, amax is enforced to be a scalar.
+            scale = 448.0 / amax
+            outputs = cuda_ext.fake_e4m3fy(inputs * scale) / scale
+
+        # Zero out values that are tiny.
+        # Tiny values could lead to tiny amax and then large scale which cause overflow/saturation
+        # and won't go back to normal value after dividing by scale. The right behavior is to mark them
+        # as zero which also get rid of inf/nan
+        outputs[zero_mask] = 0.
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        inputs, = ctx.saved_tensors
+        amax = torch.tensor(ctx.amax if ctx.amax is not None else 448.0, dtype=torch.float32, device=inputs.device)
+        grad_inputs = _fake_tensor_quant_backward(inputs, amax, grad_outputs)
+        return grad_inputs, None, None, None
+
+
+class TensorQuantFunction(Function):
+    """A universal tensor quantization function
+
+    Take an input tensor, output an quantized tensor. The granularity of scale can be interpreted from the
+    shape of amax.
+    output_dtype indicates whether the quantized value will be stored in integer or float. The reason we want to store
+    it in float is the pytorch function takes the quantized value may not accept integer input, e.g. Conv2D.
+
+    It uses 2^num_bits -1 values instead of 2^num_bits. e.g., for num_bits=8, it uses [-127, 127] instead of [-128, 127]
+    """
+
+    @staticmethod
+    @symbolic_helper.parse_args("v", "t", "i", "b", "b")
+    def symbolic(g, inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+        return _onnx_int8_helper(g, inputs, amax, num_bits, unsigned, narrow_range)
+
+    @staticmethod
+    def forward(ctx, inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+        """
+
+        Follow tensorflow convention, max value is passed in and used to decide scale, instead of inputing scale
+        directly. Though inputing scale directly may be more natural to use.
+
+        Args:
+            ctx: A Context object to store tensors for backward.
+            inputs: A Tensor of type float32.
+            amax: A Tensor of type float32. Inputs will be quantized within range [-amax, amax]
+                amax will be broadcasted to inputs tensor.
+            num_bits: A integer used to calculate scaling factor, scale = (2^(num_bits-1) - 1) / max
+                Effectively, it indicates how many integer bits is used to represent the value. Default 8.
+            output_dtype: A type of Tensor. torch.int32 or torch.float32.
+            unsigned: A boolean. Use unsigned integer range. E.g. [0, 255] for num_bits=8. Default False.
+            narrow_range: A boolean. Use symmetric integer range for signed quantization
+                E.g. [-127,127] instead of [-128,127] for num_bits=8. Default True.
+
+        Returns:
+            outputs: A Tensor of type output_dtype.
+            scale: A Tensor of type float32. outputs / scale will dequantize outputs tensor.
+
+        Raises:
+            ValueError:
+        """
+        ctx.save_for_backward(inputs, amax)
+        outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+        # Check if scale overflows FP16
+        if outputs.dtype == torch.half and scale.max() > 65504:
+            raise ValueError("scale is too large for FP16 with amax={}".format(amax))
+        return outputs, scale.to(inputs.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_outputs, grad_scale):
+        """
+        Implements straight through estimation with clipping. For -amax <= input <= amax
+        the gradient passes straight through, otherwise the gradient is zero.
+
+        Args:
+            ctx: A Context object with saved tensors from forward.
+            grad_outputs: A tensor of gradient of outputs.
+            grad_scale: A tensor of gradient of scale.
+
+        Returns:
+            grad_inputs: A tensor of gradient.
+        """
+        inputs, amax = ctx.saved_tensors
+        zero = grad_outputs.new_zeros(1)  # create a zero tensor with the same type and device
+        grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero)
+        return grad_inputs, None, None, None, None
+
+
+class LegacyFakeTensorQuantFunction(Function):
+    """Fake version of TensorQuantFunction
+    See comments of TensorQuantFunction, arguments are the same.
+    """
+
+    @staticmethod
+    def forward(ctx, inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+        ctx.save_for_backward(inputs, amax)
+        outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+        return outputs / scale.to(inputs.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        inputs, amax = ctx.saved_tensors
+        zero = grad_outputs.new_zeros(1)
+        grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero)
+        return grad_inputs, None, None, None, None
+
+
+def _tensor_quant(inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+    """Shared function body between TensorQuantFunction and FakeTensorQuantFunction"""
+    # Fine scale, per channel scale will be handled by broadcasting, which could be tricky. Pop a warning.
+    if isinstance(amax, torch.Tensor) and inputs.dim() != amax.dim():
+        logging.debug("amax %s has different shape than inputs %s. Make sure broadcast works as expected!", amax.size(),
+                      inputs.size())
+
+    logging.debug("{} bits quantization on shape {} tensor.".format(num_bits, inputs.size()))
+
+    if unsigned:
+        if inputs.min() < 0.:
+            raise TypeError("Negative values encountered in unsigned quantization.")
+
+    # Computation must be in FP32 to prevent potential over flow.
+    input_dtype = inputs.dtype
+    if inputs.dtype == torch.half:
+        inputs = inputs.float()
+    if amax.dtype == torch.half:
+        amax = amax.float()
+
+    min_amax = amax.min()
+    if min_amax < 0:
+        raise ValueError("Negative values in amax")
+
+    max_bound = torch.tensor((2.0**(num_bits - 1 + int(unsigned))) - 1.0, device=amax.device)
+    if unsigned:
+        min_bound = 0
+    elif narrow_range:
+        min_bound = -max_bound
+    else:
+        min_bound = -max_bound - 1
+    scale = max_bound / amax
+
+    epsilon = 1. / (1 << 24)
+    if min_amax <= epsilon:  # Treat amax smaller than minimum representable of fp16 0
+        zero_amax_mask = (amax <= epsilon)
+        scale[zero_amax_mask] = 0  # Value quantized with amax=0 should all be 0
+
+    outputs = torch.clamp((inputs * scale).round_(), min_bound, max_bound)
+
+    if min_amax <= epsilon:
+        scale[zero_amax_mask] = 1.  # Return 1 makes more sense for values quantized to 0 with amax=0
+
+    if input_dtype == torch.half:
+        outputs = outputs.half()
+
+    return outputs, scale
+
+
+class FakeAffineTensorQuantFunction(Function):
+    """Fake version of affine quantization
+
+    gemmlowp style scale+shift quantization. See more details in
+    https://github.com/google/gemmlowp/blob/master/doc/quantization.md.
+
+    We DO NOT recommend affine quantization on weights for performance reason. There might be value to affine quantize
+    activation as it can be cancelled by bias and comes with no performance penalty. This functionality is only added
+    for experimental purpose.
+    """
+
+    @staticmethod
+    def forward(ctx, inputs, min_range, max_range, num_bits=8):
+        """
+
+        As it will be only applied on activation with per tensor granularity, broadcast is not needed.
+
+        Args:
+            ctx: Pytorch convention.
+            inputs: A Tensor of type float32.
+            min_range: A float.
+            max_range: A float.
+            num_bits: An integer
+
+        Returns:
+            outputs: A Tensor of type output_dtype
+        """
+        logging.debug("{} bits quantization on shape {} tensor.".format(num_bits, inputs.size()))
+        ctx.save_for_backward(inputs, min_range, max_range)
+
+        step_size = (max_range - min_range) / (2.0**num_bits - 1)
+
+        min_bound = -2.0**(num_bits - 1)
+        max_bound = 2.0**(num_bits - 1) - 1
+
+        quant_zero = torch.round(min_range / step_size) - min_bound
+        quantized = torch.round(inputs / step_size) - quant_zero
+        quantized = torch.clamp(quantized, min_bound, max_bound)
+
+        outputs = (quantized + quant_zero) * step_size
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        """
+        Args:
+            ctx: Pytorch convention.
+            grad_output: A tensor of gradient of outputs
+
+        Returns:
+            grad_inputs: A tensor of gradient
+        """
+
+        inputs, min_range, max_range = ctx.saved_tensors
+        zero = grad_outputs.new_zeros(1)
+        grad_inputs = torch.where((inputs <= max_range) * (inputs >= min_range), grad_outputs, zero)
+        return grad_inputs, None, None, None
+
+
+tensor_quant = TensorQuantFunction.apply
+legacy_fake_tensor_quant = LegacyFakeTensorQuantFunction.apply
+fake_tensor_quant = FakeTensorQuantFunction.apply
+fake_affine_tensor_quant = FakeAffineTensorQuantFunction.apply
+scaled_e4m3 = ScaledE4M3Function.apply

tools/pytorch-quantization/pytorch_quantization/utils/__init__.py

@@ -0,0 +1,21 @@
+#
+# 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.
+#
+
+
+"""Main entry of all utils"""
+
+from .reduce_amax import reduce_amax

tools/pytorch-quantization/pytorch_quantization/utils/quant_logging.py

@@ -0,0 +1,27 @@
+#
+# 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.
+#
+
+
+"""A WAR for codes that messes up logging format"""
+
+import logging
+
+def reset_logger_handler():
+    """Remove all handler in root logger"""
+    root_logger = logging.getLogger()
+    while root_logger.handlers:
+        root_logger.removeHandler(root_logger.handlers[0])

tools/pytorch-quantization/pytorch_quantization/utils/reduce_amax.py

@@ -0,0 +1,63 @@
+#
+# 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.
+#
+
+
+"""Function to get absolute maximum of a tensor
+Follow numpy fashion, which is more generic as pytorch's
+"""
+
+import torch
+
+def reduce_amax(input, axis=None, keepdims=True):
+    """Compute the absolute maximum value of a tensor.
+
+    Reduces input_tensor along the dimensions given in axis. Unless keepdims is true,
+    the rank of the tensor is reduced by 1 for each entry in axis. If keepdims is true,
+    the reduced dimensions are retained with length 1.
+
+    .. note::
+        Gradient computeation is disabled as this function is never meant learning reduces amax
+
+    Args:
+        input: Input tensor
+        axis: The dimensions to reduce. None or int or tuple of ints. If None (the default),
+            reduces all dimensions. Must be in the range [-rank(input_tensor), rank(input_tensor)).
+        keepdims: A boolean. If true, retains reduced dimensions with length 1. Default True
+        granularity: DEPRECTED. specifies if the statistic has to be calculated at tensor or channel granularity
+
+    Returns:
+        The reduced tensor.
+
+    Raises:
+        ValueError: Any axis which doesn't make sense or is not supported
+        ValueError: If unknown granularity is passed in.
+    """
+    with torch.no_grad():
+        output = input.abs()
+        if axis is None:
+            output = torch.max(output)
+        else:
+            if isinstance(axis, int):
+                output, _ = torch.max(output, dim=axis, keepdim=keepdims)
+            else:
+                if isinstance(axis, tuple) and len(axis) > input.dim():
+                    raise ValueError("Cannot reduce more axes than tensor's dim.")
+                for i in axis:
+                    output, _ = torch.max(output, dim=i, keepdim=True)
+                if not keepdims or output.numel() == 1:
+                    output.squeeze_()
+        return output