tinynn/graph/quantization/quantizer.py

import copy
import functools
import inspect
import os
import queue
import re
import sys
import typing

from distutils.version import LooseVersion

import torch
import torch.nn as nn
import torch.nn.intrinsic as nni
import torch.nn.quantized as nnq
import torch.nn.quantized.dynamic as nnqd
import torch.quantization as torch_q

from torch.nn.parallel.data_parallel import DataParallel
from torch.nn.parallel.distributed import DistributedDataParallel

from tinynn.graph.quantization.fake_quantize import (
    FakeQuantizeBFloat16,
    FakeQuantizeTFLite,
)
from tinynn.graph.quantization.modules import QGLU, QHardsigmoid, QPReLU, QSiLU, QLayerNorm, QRMSNorm
from tinynn.graph.quantization.observer import (
    HistogramObserverKL,
    MinMaxObserver,
    PerChannelMinMaxObserver,
)
from tinynn.graph.quantization.qat_modules import (
    Conv1d,
    ConvTranspose1d,
    ConvTranspose2d,
    ConvTransposeBn2d,
)

if LooseVersion(torch.__version__) >= LooseVersion("1.13.0"):
    from tinynn.graph.quantization.quantizable.gru import GRU as QuantizableGRU
    from torch.ao.nn.quantizable.modules.rnn import LSTM as QuantizableLSTM
else:
    QuantizableGRU = None
    QuantizableLSTM = None

from tinynn.graph.tracer import (
    ConstantNode,
    TraceFunction,
    TraceGraph,
    TraceNode,
    load_creation_funcs,
    module_constructor_lines,
    override_current_trace_graph,
    qualified_name,
    trace,
)
from tinynn.util.train_util import get_logger, get_module_device
from tinynn.util.util import import_from_path
from tinynn.graph.quantization.utils import fake_quantize

from . import fused_modules as fm

try:
    import ruamel_yaml as yaml
    from ruamel_yaml import CommentedMap
except ModuleNotFoundError:
    import ruamel.yaml as yaml
    from ruamel.yaml import CommentedMap


# Fusable OPs for Quantize Aware Training
FUSE_RULE_LIST = {
    (torch.nn.Conv1d, torch.nn.BatchNorm1d),
    (torch.nn.Conv1d, torch.nn.BatchNorm1d, torch.nn.ReLU),
    (torch.nn.Conv2d, torch.nn.BatchNorm2d),
    (torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.ReLU),
    (torch.nn.Conv3d, torch.nn.BatchNorm3d),
    (torch.nn.Conv3d, torch.nn.BatchNorm3d, torch.nn.ReLU),
    (torch.nn.Conv1d, torch.nn.ReLU),
    (torch.nn.Conv2d, torch.nn.ReLU),
    (torch.nn.Conv3d, torch.nn.ReLU),
    (torch.nn.Linear, torch.nn.ReLU),
    (torch.nn.BatchNorm2d, torch.nn.ReLU),
    (torch.nn.BatchNorm3d, torch.nn.ReLU),
    (torch.nn.ConvTranspose2d, torch.nn.BatchNorm2d),
}

FUSE_RULE_LIST_PTQ_ONLY = {
    (nn.Linear, nn.BatchNorm1d): '1.8.0',
    (nn.ConvTranspose1d, nn.BatchNorm1d): '1.11.0',
    (nn.ConvTranspose3d, nn.BatchNorm3d): '1.11.0',
    (nn.BatchNorm2d, nn.Conv2d): None,
    (nn.BatchNorm2d, nn.Conv2d, nn.ReLU): None,
}

FUSE_RULE_LIST_EXTRA = {
    (torch.nn.Conv1d, torch.nn.BatchNorm1d, torch.nn.ReLU6),
    (torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.ReLU6),
    (torch.nn.Conv2d, torch.nn.BatchNorm2d, 'clamp_with_fusion'),
    (torch.nn.Conv3d, torch.nn.BatchNorm3d, torch.nn.ReLU6),
    (torch.nn.Conv1d, torch.nn.ReLU6),
    (torch.nn.Conv2d, torch.nn.ReLU6),
    (torch.nn.Conv2d, 'clamp_with_fusion'),
    (torch.nn.Conv3d, torch.nn.ReLU6),
    (torch.nn.Linear, torch.nn.ReLU6),
    (torch.nn.Linear, 'clamp_with_fusion'),
    (torch.nn.Linear, torch.nn.BatchNorm1d, torch.nn.ReLU6),
    (torch.nn.Linear, torch.nn.BatchNorm1d, 'clamp_with_fusion'),
    (torch.nn.BatchNorm2d, torch.nn.ReLU6),
    (torch.nn.BatchNorm2d, 'clamp_with_fusion'),
    (torch.nn.BatchNorm3d, torch.nn.ReLU6),
    ('add', torch.nn.ReLU6),
    ('add', 'relu6'),
    ('add', 'clamp_with_fusion'),
    (torch.nn.ConvTranspose2d, torch.nn.ReLU),
    (torch.nn.ConvTranspose2d, torch.nn.ReLU6),
    (torch.nn.ConvTranspose2d, 'clamp_with_fusion'),
    (torch.nn.ConvTranspose2d, torch.nn.BatchNorm2d, torch.nn.ReLU),
    (torch.nn.ConvTranspose2d, torch.nn.BatchNorm2d, torch.nn.ReLU6),
    (torch.nn.ConvTranspose2d, torch.nn.BatchNorm2d, 'clamp_with_fusion'),
}

FUSE_FALLBACK_DICT = {'clamp_with_fusion': 'clamp'}

FUSE_QAT_MODULES = {
    nn.Conv1d: Conv1d,
    nn.ConvTranspose1d: ConvTranspose1d,
    nn.ConvTranspose2d: ConvTranspose2d,
    fm.ConvTransposeBn2d: ConvTransposeBn2d,
}

FUSE_QAT_MODULES_CUSTOM = {}
FUSE_PTQ_MODULES_CUSTOM = {}

if LooseVersion(torch.__version__) >= '1.13.0':
    from .quantizable.gru import GRU

    FUSE_QAT_MODULES_CUSTOM.update({nn.GRU: GRU})
    FUSE_PTQ_MODULES_CUSTOM.update({nn.GRU: GRU})

FUSE_QAT_MODULES_CVT = {Conv1d: nnq.Conv1d}
if hasattr(nnq, 'ConvTranspose1d'):
    FUSE_QAT_MODULES_CVT.update(
        {
            ConvTranspose1d: nnq.ConvTranspose1d,
            ConvTranspose2d: nnq.ConvTranspose2d,
        }
    )

REWRITE_TO_FUSE_RULE_LIST = {
    (torch.nn.Linear, torch.nn.BatchNorm1d),
}

KNOWN_QSTATS = {
    nn.Softmax: (0, 256.0),
    'softmax': (0, 256.0),
    nn.LogSoftmax: (255, 16.0),
    'log_softmax': (255, 16.0),
}

REWRITE_QUANTIZABLE_RULE_LIST = {
    ('abs',),
    ('sum',),
    ('bmm',),
    ('matmul',),
    ('truediv',),
    ('clamp_min',),
    ('clamp_max',),
    ('sqrt', 'reciprocal'),
}

UNSUPPORTED_PYTORCH_QUANTIZATION_OP_LIST = {
    'pad': '1.7.0',
    'pow': None,
    'truediv': None,
    'sqrt': None,
    'rsqrt': None,
    'atan2': None,
    'atan': None,
    'sin': None,
    'cos': None,
    'hardsigmoid': None,
    'silu': None,
    'reciprocal': None,
    'exp': None,
    'layer_norm': None,
    'instance_norm': None,
    'softmax': None,
    'log_softmax': None,
    'mm': None,
    'matmul': None,
    'bmm': None,
    'abs': None,
    'sum': None,
    'clamp_min': None,
    'clamp_max': None,
    'prelu': None,
    'elu': None,
    'glu': None,
    'group_norm': None,
    'log': None,
    'std': None,
    'var': None,
    'norm': None,
    nn.LSTM: '1.13.0',
    nn.ConvTranspose2d: '1.7.0',
    nn.ConstantPad1d: '1.7.0',
    nn.ConstantPad2d: '1.7.0',
    nn.ConstantPad3d: '1.7.0',
    nn.ZeroPad2d: '1.7.0',
    nn.RNN: None,
    nn.GRU: '1.13.0',
    nn.LayerNorm: None,
    nn.InstanceNorm1d: None,
    nn.InstanceNorm2d: None,
    nn.Hardsigmoid: None,
    nn.Softmax: None,
    nn.LogSoftmax: None,
    nn.GLU: None,
    nn.PReLU: None,
    nn.GroupNorm: None,
}

Q_MODULES_MAPPING = {
    nn.PReLU: QPReLU,
    nn.GLU: QGLU,
    nn.Hardsigmoid: QHardsigmoid,
}

FUNCTIONAL_MODULE_MAPPING = {
    'relu': nn.ReLU,
    'relu6': nn.ReLU6,
    'leaky_relu': nn.LeakyReLU,
    'elu': nn.ELU,
    'prelu': nn.PReLU,
    'glu': nn.GLU,
    'sigmoid': nn.Sigmoid,
    'tanh': nn.Tanh,
    'hardswish': nn.Hardswish,
}

TENSORRT_OBSERVED_NODES = [
    (torch.nn, "Conv1d"),
    (torch.nn, "Conv2d"),
    (torch.nn, "Conv3d"),
    (torch.nn, "ConvTranspose1d"),
    (torch.nn, "ConvTranspose2d"),
    (torch.nn, "ConvTranspose3d"),
    (torch.nn, "Linear"),
    (torch.nn, "LSTM"),
    (torch.nn, "LSTMCell"),
    (torch.nn, "AvgPool1d"),
    (torch.nn, "AvgPool2d"),
    (torch.nn, "AvgPool3d"),
    (torch.nn, "AdaptiveAvgPool1d"),
    (torch.nn, "AdaptiveAvgPool2d"),
    (torch.nn, "AdaptiveAvgPool3d"),
    'avg_pool2d',
    'adaptive_avg_pool2d',
]

if hasattr(nn, 'SiLU'):
    UNSUPPORTED_PYTORCH_QUANTIZATION_OP_LIST.update({nn.SiLU: None})
    Q_MODULES_MAPPING.update({nn.SiLU: QSiLU})
    FUNCTIONAL_MODULE_MAPPING.update({'silu': nn.SiLU})

if hasattr(nn, 'LayerNorm'):
    Q_MODULES_MAPPING.update({nn.LayerNorm: QLayerNorm})

if hasattr(nn, 'RMSNorm'):
    Q_MODULES_MAPPING.update({nn.RMSNorm: QRMSNorm})

# Processed QAT fuse rules
processed_qat_rules = {}
processed_ptq_rules = {}
processed_extra_q_rules = {}
processed_all_qat_rules = {}
processed_all_ptq_rules = {}

# Constant func names
creation_func_names = []

# Processed rewrite rules for fusing
processed_rewrite_to_fuse_rules = {}

log = get_logger(__name__, 'WARNING')


class QATQuantizer(object):
    rewrite_graph: bool
    force_overwrite: bool
    is_input_quantized: typing.Optional[typing.Tuple[bool]]
    quantized_input_stats: typing.Optional[typing.List[typing.Optional[typing.Tuple[float, float]]]]
    quantized_op_stats: typing.Optional[typing.List[typing.Optional[typing.Tuple[float, float]]]]
    set_quantizable_op_stats: bool
    backend: str
    remove_weights_after_load: bool
    asymmetric: bool
    per_tensor: bool
    disable_requantization_for_cat: bool
    dynamic_lstm_quant: bool
    rounding_mode: str
    leaf_nodes: typing.Optional[typing.List[nn.Module]]
    swap_nodes: typing.Optional[typing.List[typing.Tuple[nn.Module, nn.Module]]]
    legacy_fq: bool
    extra_tracer_opts: typing.Optional[typing.Dict]
    layerwise_config: typing.Dict[str, bool]
    layerwise_default: bool
    override_qconfig_func: typing.Optional[typing.Callable[[str, nn.Module], typing.Optional[torch_q.QConfig]]]
    effective_layers: typing.List[str]
    ignore_layerwise_config: bool
    fused_layerwise_config: bool
    inplace: bool
    train_mode_dict: typing.Dict[nn.Module, bool]

    def __init__(self, model, dummy_input, work_dir: typing.Optional[str] = None, config: typing.Optional[dict] = None):
        """ Constructs a new QATQuantizer object

        Args:
            model: The model to be quantized
            dummy_input: A viable input to the model
            work_dir (typing.Optional[str], optional): The working directory in which the intermediate files will be \
                generated. Defaults to None, in which case "output" will be used.
            config (typing.Optional[dict]): Options for the quantizer
        """

        super().__init__()

        if isinstance(model, DataParallel) or isinstance(model, DistributedDataParallel):
            self.model = model.module
        else:
            self.model = model

        self.dummy_input = dummy_input
        self.work_dir = 'out' if work_dir is None else work_dir

        self.parse_config(config)

        if sys.platform == 'win32' and self.backend == 'qnnpack':
            log.error('Quantization backend qnnpack is likely unsupported on Windows. Please use fbgemm instead.')

        if self.backend not in ('fbgemm', 'qnnpack', 'onnx', 'tensorrt'):
            log.warning(f'Quantization backend {self.backend} is not tested. Please use at your risk.')

        if self.backend == 'fbgemm':
            assert self.asymmetric, "Symmetric quantizaton for FBGEMM not supported"
            assert (
                not self.per_tensor
            ), "Per-tensor quantizaton for FBGEMM not supported, please use per-channel quantization instead"

        if self.backend == 'tensorrt':
            if self.asymmetric:
                log.warning('Asymmetric quantizaton for TensorRT not supported')

        if self.disable_requantization_for_cat is None:
            if not self.per_tensor:
                self.disable_requantization_for_cat = True
            else:
                self.disable_requantization_for_cat = False

        self.extra_qparams_mappings = []

        assert (
            self.per_tensor or self.disable_requantization_for_cat
        ), "`disable_requantization_for_cat=True` is required for per-channel quantization"

        if self.legacy_fq:
            version = None
            if type(self).__name__ == 'QATQuantizer':
                version = '1.10.0'
            elif type(self).__name__ == 'PostQuantizer':
                version = '1.12.0'

            if version is None or LooseVersion(torch.__version__) < version:
                log.info(f'legacy_fq=True is only available for QATQuantizer and PostQuantizer with PyTorch {version}+')
                self.legacy_fq = False

        self.leaf_nodes = None
        self.swap_nodes = None

        self.train_mode_dict = {}

        self.layerwise_config = CommentedMap()
        self.effective_layers = []
        self.layerwise_default = True
        if config is not None and 'layerwise_config' in config:
            self.layerwise_config.update(config['layerwise_config'])

        self.lstm_origin_weight_dict = {}

    def parse_config(self, config: typing.Optional[dict]):
        default_values = {
            'rewrite_graph': True,
            'force_overwrite': True,
            'is_input_quantized': None,
            'backend': 'qnnpack',
            'remove_weights_after_load': False,
            'asymmetric': True,
            'per_tensor': True,
            'disable_requantization_for_cat': None,
            'quantized_input_stats': None,
            'dynamic_lstm_quant': False,
            'quantized_op_stats': None,
            'set_quantizable_op_stats': False,
            'rounding_mode': 'pytorch',
            'algorithm': 'l2',
            'fuse_only': False,
            'legacy_fq': True,
            'extra_tracer_opts': {},
            'fused_layerwise_config': True,
            'ignore_layerwise_config': False,
            'inplace': False,
            'override_qconfig_func': None,
            'quantize_op_action': {},
        }

        if config is None:
            config = dict()

        for k, v in default_values.items():
            actual_v = config.get(k, v)
            setattr(self, k, actual_v)

    def set_model_mode(self, model: nn.Module = None, preserve: bool = False):
        """Sets the mode (train/eval) of the model

        Args:
            model (nn.Module, optional): A PyTorch model. Defaults to None.
            preserve (bool, optional): Whether the old state is preserved. Defaults to False.
        """
        if model is None:
            model = self.model

        if preserve:
            for m in model.modules():
                training = getattr(m, 'training', None)
                if isinstance(training, bool):
                    self.train_mode_dict[m] = training

        model.train()

    def restore_model_mode(self, model: nn.Module = None):
        """Restores the mode (train/eval) of the model

        Args:
            model (nn.Module, optional): A PyTorch model. Defaults to None.
        """
        if model is None:
            model = self.model

        for m, training in self.train_mode_dict.items():
            if training:
                m.train()
            else:
                m.eval()

        self.train_mode_dict.clear()

    def quantize(self) -> nn.Module:
        """Performs QAT rewrite and preparation

        Returns:
            nn.Module: The QAT-ready model
        """

        # We need a model in training mode so that QAT could take place
        self.set_model_mode(preserve=True)

        if self.rewrite_graph:
            # Retrives the name of the model from type info
            model_name = type(self.model).__name__
            model_name_q = f'Q{model_name}'
            model_name_q_lower = f'{model_name}_q'.lower()
            model_ns = f'tinynn_rewritten_models.{model_name_q}'

            model_code_path = os.path.join(self.work_dir, f'{model_name_q_lower}.py')
            model_weights_path = os.path.join(self.work_dir, f'{model_name_q_lower}.pth')

            config_path = os.path.join(self.work_dir, f'{model_name_q_lower}_config.yml')

            yaml_ = yaml.YAML()
            if len(self.layerwise_config) == 0 and not self.ignore_layerwise_config:
                if os.path.exists(config_path):
                    with open(config_path, 'r') as f:
                        self.layerwise_config = yaml_.load(f)

            if 'default' in self.layerwise_config:
                self.layerwise_default = self.layerwise_config['default']
                del self.layerwise_config['default']

            # Generate the code for the modified model
            # We will try to do some op fusion and rewriting in the `rewrite_quantize_graph` function.
            # By default, we will try to insert QuantStubs before every input and DeQuantStubs after every output in
            # the generated graph. If this doesn't suit your needs, e.g. you have intergal/quantized inputs or want to
            # skip the quantization rewrite for some ops, then you may modify the code generated freely and remember to
            # skip this step so it won't be overwritten.
            if self.force_overwrite or not os.path.exists(model_code_path) or not os.path.exists(model_weights_path):
                # After tracing the model, we will get a TraceGraph object
                graph = trace(self.model, self.dummy_input, **self.extra_tracer_opts)
                self.rewrite_quantize_graph(graph)
                if self.inplace:
                    graph.inplace_commit()
                else:
                    graph.generate_code(model_code_path, model_weights_path, model_name_q)

            os.makedirs(self.work_dir, exist_ok=True)
            with open(config_path, 'w') as f:
                yaml_.dump(self.layerwise_config, f)

            if self.inplace:
                rewritten_model = self.model
            else:
                # Import the new model
                rewritten_model = import_from_path(model_ns, model_code_path, model_name_q)()
                rewritten_model.load_state_dict(torch.load(model_weights_path))

                device = get_module_device(self.model)
                if device is not None:
                    rewritten_model.to(device=device)

                # Remove the weights file to save space
                if self.remove_weights_after_load:
                    os.unlink(model_weights_path)

                # Restores the mode of the model
                self.restore_model_mode()

                # Set the model to training mode before tracing again, since we need to perform QAT
                self.set_model_mode(rewritten_model)

            rewritten_graph = trace(rewritten_model, self.dummy_input)

            # Update the model so that the original one can be released
            self.model = rewritten_model
        else:
            # After tracing the model, we will get a TraceGraph object
            graph = trace(self.model, self.dummy_input, **self.extra_tracer_opts)
            rewritten_graph = graph

        # Fuse the modules (e.g conv, bn, relu) in the computation graph according to the fuse rules.
        # By default, we assume all input tensors are of floating type.
        # If you want to use quantized/integral inputs, then you may need to pass in `is_input_quantized`.
        qat_model = self.prepare_qat(rewritten_graph, self.is_input_quantized, self.backend, self.fuse_only)

        return qat_model

    def prepare_qat_prep(
        self,
        graph: TraceGraph,
        is_input_quantized: typing.Optional[typing.Tuple[bool]] = None,
        backend: str = 'qnnpack',
    ):
        """Some common logic before calling torch.quantization.prepare[_qat]

        Args:
            graph (TraceGraph): The computation graph of the model
            is_input_quantized (typing.Union[typing.Tuple[bool]], optional): Whether the input tensor(s) is (are) \
                quantized. Defaults to None.
            backend (str, optional): The backend of quantization. Defaults to 'qnnpack'.

        """

        qat_analysis_queue = queue.Queue()
        visited = set()

        def _qat_analysis(node: TraceNode, quantized: bool):
            # Find quantized subgraphs in the whole computation graph

            if node.unique_name in visited:
                return

            visited.add(node.unique_name)

            if node in graph.output_nodes:
                return

            # TODO: Enable QAT analysis for TensorRT
            if self.backend == 'tensorrt':
                quantized = True
                node.quantized = True
            elif type(node.module) is torch_q.QuantStub:
                quantized = True
                node.quantized = quantized
            elif type(node.module) is torch_q.DeQuantStub:
                node.quantized = True
                quantized = False
            else:
                node.quantized = quantized
                log.debug(f"[QUANTIZED]{node.unique_name}:{quantized}")

            for i, n in enumerate(node.next_nodes):
                if type(n.module) is TraceFunction:
                    if n.kind() in ('shape', 'size', 'dtype', 'device'):
                        continue
                    if n.kind() == 'expand_as' and i > 0:
                        continue
                qat_analysis_queue.put((n, quantized))

        if is_input_quantized is not None:
            assert len(is_input_quantized) == len(graph.input_nodes)

            for n, q in zip(graph.input_nodes, is_input_quantized):
                qat_analysis_queue.put((n, q))
        else:
            for n in graph.input_nodes:
                qat_analysis_queue.put((n, False))

        creation_func_names = load_creation_func_names()

        def _is_extra_constant_nodes(node, custom_data):
            return node.full_name() in creation_func_names

        extra_constant_nodes = graph.filter_forward_nodes(_is_extra_constant_nodes)

        def _is_params_in_module(node, custom_data):
            if len(node.prev_nodes) == 1 and len(node.next_nodes) == 1:
                if len(node.prev_tensors) == 1 and len(node.next_tensors) == 1:
                    if isinstance(node.prev_tensors[0], nn.Module) and not isinstance(
                        node.prev_tensors[0], nnq.FloatFunctional
                    ):
                        return True
            return False

        param_nodes = graph.filter_forward_nodes(_is_params_in_module)

        for n in graph.constant_nodes + extra_constant_nodes:
            qat_analysis_queue.put((n, not torch.is_floating_point(n.next_tensors[0])))

        while not qat_analysis_queue.empty():
            node, quantized = qat_analysis_queue.get()
            if not graph.quantized:
                graph.quantized = graph.quantized or quantized
            _qat_analysis(node, quantized)

        q_dict = {}
        for n in graph.forward_nodes:
            if isinstance(n.module, nn.Module):
                q_dict[n.module] = q_dict.get(n.module, False) or n.quantized

        for n in graph.forward_nodes:
            if n.module in q_dict:
                n.quantized = q_dict[n.module]

        for n in param_nodes:
            prev_node = n.prev_nodes[0]
            next_node = n.next_nodes[0]
            is_known_mod = prev_node.kind().__name__ in (
                'Conv1d',
                'Conv2d',
                'Linear',
                'ConvTranspose1d',
                'ConvTranspose2d',
            )
            if is_known_mod and n.module.full_name == 'weight' and prev_node.quantized:
                if next_node.type() is torch_q.QuantStub:
                    mod = nn.Sequential(torch_q.DeQuantStub(), torch_q.QuantStub())
                    orig_mod = next_node.module
                    next_node.module = mod
                    setattr(graph.module, next_node.original_name, next_node.module)
                    graph.module_original_name_dict[id(mod)] = graph.module_original_name_dict[id(orig_mod)]
                    graph.module_unique_name_dict[id(mod)] = graph.module_unique_name_dict[id(orig_mod)]
                continue
            qat_analysis_queue.put((n, not torch.is_floating_point(n.next_tensors[0])))

        while not qat_analysis_queue.empty():
            node, quantized = qat_analysis_queue.get()
            if not graph.quantized:
                graph.quantized = graph.quantized or quantized
            _qat_analysis(node, quantized)

        log.debug("qat analysis over")

        if not graph.quantized:
            return

        if isinstance(self, PostQuantizer):
            processed_rules = load_processed_ptq_rules()
        else:
            processed_rules = load_processed_qat_rules()

        is_fusable = functools.partial(self.is_fusable, current_rules=processed_rules, graph=graph)

        def _find_quantized_module_nodes(node: TraceNode, custom_node):
            # Find quantized module nodes
            return node.type() in Q_MODULES_MAPPING and node.quantized

        # Replace module nodes with our custom variants
        quantized_mod_nodes = graph.filter_forward_nodes(_find_quantized_module_nodes)

        type_dict = {}
        for node in quantized_mod_nodes:
            node_type = node.type()
            type_dict.setdefault(node_type, [])
            type_dict[node_type].append(node)

        for node_type, nodes in type_dict.items():
            graph.update_submodule_in_nodes_from_predicate(nodes, Q_MODULES_MAPPING[node_type], self.inplace)

        custom_data = ([], set())
        graph.filter_forward_nodes(is_fusable, custom_data, reverse=True)
        quant_list = custom_data[0]
        log.info(f'found nodes to fuse: {quant_list}')

        new_fuser_func = fm.gen_fuse_known_modules_wrapper(
            sys.modules['torch.quantization.fuse_modules'].fuse_known_modules
        )

        if self.backend != 'tensorrt':
            is_qat = type(self) is QATQuantizer
            for quant_nodes in quant_list:
                if self.inplace:
                    quant_nodes = [re.sub('get_submodule\\("(.*?)"\\)', '\\1', x) for x in quant_nodes]
                    quant_nodes = [re.sub('\\[("|)(.*?)("|)\\]', '.\\2', x) for x in quant_nodes]

                if LooseVersion(torch.__version__) >= '1.11.0' and LooseVersion(torch.__version__) < '1.14.0':
                    # See https://github.com/pytorch/pytorch/pull/88193
                    sys.modules['torch.quantization.fuse_modules']._fuse_modules(
                        graph.module, quant_nodes, is_qat=is_qat, inplace=True, fuser_func=new_fuser_func
                    )
                elif is_qat and LooseVersion(torch.__version__) >= '1.14.0':
                    # See https://github.com/pytorch/pytorch/issues/74028
                    torch.ao.quantization.fuse_modules_qat(
                        graph.module, quant_nodes, fuser_func=new_fuser_func, inplace=True
                    )
                else:
                    torch_q.fuse_modules(graph.module, quant_nodes, fuser_func=new_fuser_func, inplace=True)

        self.prepare_qconfig(graph, backend)
        self.override_qconfig(graph.module)

    def override_qconfig(self, module: nn.Module):
        """Sets qconfig according to the provided `override_qconfig_func`

        Args:
            module (nn.Module): The model to be prepared for quantization
        """
        if self.override_qconfig_func is not None:
            q = queue.Queue()
            q.put(('', module))

            while not q.empty():
                n, m = q.get()

                new_qconfig = self.override_qconfig_func(n, m)
                if new_qconfig is not None:
                    m.qconfig = new_qconfig
                else:
                    for cn, cm in m.named_children():
                        fn = cn
                        if len(n) > 0:
                            fn = f'{n}.{fn}'
                        q.put((fn, cm))

    def prepare_qconfig(self, graph: TraceGraph, backend: str):
        """Prepare qconfig for various configurations.

        Args:
            graph (TraceGraph): The computation graph of the model
            backend (str, optional): The backend of quantization
        """

        log.info('setting qat backend and call prepare_qat')
        actual_backend = backend
        if backend in ('onnx', 'tensorrt'):
            actual_backend = 'qnnpack'
        if not self.legacy_fq:
            qconfig = torch_q.get_default_qat_qconfig(actual_backend)
        else:
            if LooseVersion(torch.__version__) >= '1.13.0':
                # See https://github.com/pytorch/pytorch/pull/88876
                qconfig = torch_q.QConfig(
                    activation=torch_q.FakeQuantize.with_args(
                        observer=torch_q.MovingAverageMinMaxObserver, quant_min=0, quant_max=255, reduce_range=False
                    ),
                    weight=torch_q.default_weight_fake_quant,
                )
            else:
                version = None
                if LooseVersion(torch.__version__) >= '1.12.0':
                    version = 0
                qconfig = torch_q.get_default_qat_qconfig(actual_backend, version)

        qconfig_c = None
        if self.rounding_mode == 'tflite':
            q_a = FakeQuantizeTFLite.with_args(*qconfig.activation.p.args, **qconfig.activation.p.keywords)
            q_w = FakeQuantizeTFLite.with_args(*qconfig.weight.p.args, **qconfig.weight.p.keywords)
            qconfig = torch_q.QConfig(q_a, q_w)
        if backend == 'qnnpack':
            if not self.asymmetric:
                sym_fq = qconfig.activation.with_args(
                    observer=torch_q.MovingAverageMinMaxObserver,
                    quant_min=0,
                    quant_max=255,
                    dtype=torch.quint8,
                    qscheme=torch.per_tensor_symmetric,
                    reduce_range=False,
                )
                qconfig = torch_q.QConfig(sym_fq, qconfig.weight)
            if not self.per_tensor:
                sym_fq = qconfig.weight.with_args(
                    observer=torch_q.MovingAveragePerChannelMinMaxObserver.with_args(quant_min=-127, quant_max=127),
                    quant_min=-127,
                    quant_max=127,
                    dtype=torch.qint8,
                    qscheme=torch.per_channel_symmetric,
                    reduce_range=False,
                    ch_axis=0,
                )
                qconfig_c = torch_q.QConfig(qconfig.activation, sym_fq)
        elif backend == 'fbgemm':
            fq_type = qconfig.weight.p.func
            sym_fq = fq_type.with_args(
                observer=torch_q.MovingAverageMinMaxObserver,
                quant_min=-128,
                quant_max=127,
                dtype=torch.qint8,
                qscheme=torch.per_tensor_symmetric,
                reduce_range=False,
            )
            qconfig_c = torch_q.QConfig(qconfig.activation, sym_fq)
        elif backend in ('onnx', 'tensorrt'):
            if not self.asymmetric:
                sym_fq = qconfig.activation.with_args(
                    observer=torch_q.MovingAverageMinMaxObserver,
                    quant_min=-128,
                    quant_max=127,
                    dtype=torch.qint8,
                    qscheme=torch.per_tensor_symmetric,
                    reduce_range=False,
                )
                qconfig = torch_q.QConfig(sym_fq, qconfig.weight)
            if not self.per_tensor:
                sym_fq = qconfig.weight.with_args(
                    observer=torch_q.MovingAveragePerChannelMinMaxObserver,
                    quant_min=-128,
                    quant_max=127,
                    dtype=torch.qint8,
                    qscheme=torch.per_channel_symmetric,
                    reduce_range=False,
                    ch_axis=0,
                )
                qconfig_c = torch_q.QConfig(qconfig.activation, sym_fq)
        else:
            log.warning(f'Quantization backend {self.backend} is not tested. Please use at your risk.')

        torch.backends.quantized.engine = actual_backend
        graph.module.qconfig = qconfig
        if self.backend == 'qnnpack':
            if qconfig_c is not None:
                q = queue.Queue()
                q.put(graph.module)

                while not q.empty():
                    m = q.get()
                    if type(m).__name__ in (
                        'Conv2d',
                        'ConvBnReLU2d',
                        'ConvBn2d',
                        'ConvReLU2d',
                        'Conv1d',
                        'ConvBnReLU1d',
                        'ConvBn1d',
                    ):
                        m.qconfig = qconfig_c
                    else:
                        for c in m.children():
                            q.put(c)
        elif self.backend == 'fbgemm':
            if qconfig_c is not None:
                q = queue.Queue()
                q.put(graph.module)

                while not q.empty():
                    m = q.get()
                    if type(m).__name__ in ('Linear', 'LinearReLU'):
                        m.qconfig = qconfig_c
                    else:
                        for c in m.children():
                            q.put(c)

        def _lstm_node(node, custom_data):
            return isinstance(node.module, nn.LSTM)

        if self.dynamic_lstm_quant:
            lstm_nodes = graph.filter_forward_nodes(_lstm_node)
            for node in lstm_nodes:
                node.quantized = True
                node.module.qconfig = torch_q.default_dynamic_qconfig

    def prepare_qat(
        self,
        graph: TraceGraph,
        is_input_quantized: typing.Optional[typing.Tuple[bool]] = None,
        backend: str = 'qnnpack',
        fuse_only: bool = False,
    ) -> torch.nn.Module:
        """Prepare model for QAT training

        Args:
            graph (TraceGraph): The computation graph of the model
            is_input_quantized (typing.Union[typing.Tuple[bool]], optional): Whether the input tensor(s) is (are) \
                quantized. Defaults to None.
            backend (str, optional): The backend of quantization. Defaults to 'qnnpack'.
            fuse_only (bool, optional): Whether the returned model is only fused in PostQuantizer. Defaults to False.

        Returns:
            torch.nn.Module: The QAT-ready model
        """

        graph.module.train()

        self.prepare_qat_prep(graph, is_input_quantized, backend)

        if not graph.quantized:
            log.warning('Graph is not quantized, skip preparation')
            return graph.module

        # Unfornately, the suggested way below will try to fuse all the modules
        # even if some of the nodes are not in a quantized computation graph.
        # So we wrote some alternatives for the function.
        #   torch.quantization.prepare_qat(graph.module, inplace=True)

        if hasattr(torch_q, 'get_default_qat_module_mappings'):
            mapping = torch_q.get_default_qat_module_mappings()
        elif hasattr(torch_q, 'get_qat_module_mappings'):
            mapping = copy.deepcopy(torch_q.get_qat_module_mappings())
        else:
            mapping = copy.deepcopy(torch_q.DEFAULT_QAT_MODULE_MAPPING)

        if self.dynamic_lstm_quant:
            mapping = dict(mapping)
            mapping.update({nn.LSTM: nnqd.LSTM})

        mapping.update(FUSE_QAT_MODULES)

        q_dict = {}
        for n in graph.forward_nodes:
            if isinstance(n.module, nn.Module):
                q_dict[n.module] = q_dict.get(n.module, False) or n.quantized

        non_quantized_mods = set(m for m, s in q_dict.items() if s is False)

        if LooseVersion(torch.__version__) < LooseVersion("1.7.0"):
            model = torch_q.prepare(graph.module, inplace=True)
            for m in non_quantized_mods:
                if hasattr(m, "_forward_hooks"):
                    if len(m._forward_hooks) > 0:
                        m._forward_hooks.popitem()
                if hasattr(m, "qconfig"):
                    delattr(m, "qconfig")
                if hasattr(m, "activation_post_process"):
                    delattr(m, "activation_post_process")

            if hasattr(graph.module, 'activation_post_process'):
                if hasattr(graph.module, "_forward_hooks"):
                    if len(graph.module._forward_hooks) > 0:
                        graph.module._forward_hooks.popitem()
                delattr(graph.module, "activation_post_process")

            torch_q.convert(model, mapping, inplace=True)
        else:
            torch_q.propagate_qconfig_(graph.module, qconfig_dict=None)
            for m in non_quantized_mods:
                if hasattr(m, "qconfig"):
                    delattr(m, "qconfig")
            model = torch_q.convert(graph.module, mapping=mapping, inplace=True, remove_qconfig=False)

            if self.dynamic_lstm_quant:
                mapping.pop(nn.LSTM)

            if LooseVersion(torch.__version__) >= LooseVersion("1.13.0"):
                torch_q.propagate_qconfig_(model, qconfig_dict=None)

                for m in non_quantized_mods:
                    if hasattr(m, "qconfig"):
                        delattr(m, "qconfig")

                prepare_custom_config_dict = torch.ao.quantization.get_default_custom_config_dict()
                custom_module_class_mapping = prepare_custom_config_dict.get(
                    "float_to_observed_custom_module_class", {}
                )
                custom_module_class_mapping.update(FUSE_QAT_MODULES_CUSTOM)
                qconfig_propagation_list = torch_q.get_default_qconfig_propagation_list()

                from .quantizable import lstm, gru

                orig_from_float = torch.ao.nn.quantizable.LSTM.from_float

                torch.ao.nn.quantizable.LSTM.from_float = lstm.from_float
                GRU.from_float = gru.from_float

                def patch_observer_set(orig_func):
                    def new_no_observer_set():
                        return set(FUSE_QAT_MODULES_CUSTOM.values()) | orig_func()

                    return new_no_observer_set

                orig_no_observer_set = sys.modules['torch.ao.quantization.quantize'].no_observer_set
                sys.modules['torch.ao.quantization.quantize'].no_observer_set = patch_observer_set(orig_no_observer_set)

                if hasattr(torch_q, 'add_observer_'):
                    add_observer_func = torch_q.add_observer_
                else:
                    add_observer_func = sys.modules['torch.ao.quantization.quantize']._add_observer_

                add_observer_func(
                    model,
                    qconfig_propagation_list,
                    set(mapping.values()),
                    custom_module_class_mapping=custom_module_class_mapping,
                )

                torch.ao.nn.quantizable.LSTM.from_float = orig_from_float
                sys.modules['torch.ao.quantization.quantize'].no_observer_set = orig_no_observer_set

            else:
                torch_q.prepare(model, observer_non_leaf_module_list=set(mapping.values()), inplace=True)
            for m in non_quantized_mods:
                if hasattr(m, "_forward_hooks"):
                    if len(m._forward_hooks) > 0:
                        m._forward_hooks.popitem()
                if hasattr(m, "qconfig"):
                    delattr(m, "qconfig")
                if hasattr(m, "activation_post_process"):
                    delattr(m, "activation_post_process")

        if not self.per_tensor:
            if self.backend == 'qnnpack':
                for n, m in graph.module.named_modules():
                    if n.endswith('.weight_fake_quant'):
                        observer = getattr(m, 'activation_post_process', None)
                        if observer is not None:
                            m.quant_min = -127
                            m.quant_max = 127
                            observer.quant_min = -127
                            observer.quant_max = 127

        self.extra_qat_fusion_postprocess(graph)

        if self.disable_requantization_for_cat:
            self.disable_requantization_for_cat_pass(graph)

        if self.quantized_input_stats is not None:
            self.prepare_quantized_inputs_pass(graph)

        if self.set_quantizable_op_stats:
            if self.quantized_op_stats is None:
                self.quantized_op_stats = {}
            self.quantized_op_stats.update(KNOWN_QSTATS)

        if self.quantized_op_stats is not None:
            self.prepare_quantized_ops_pass(graph)

        if self.backend == 'tensorrt':
            q = queue.Queue()
            q.put(('', graph.module))

            while not q.empty():
                s, m = q.get()
                if isinstance(m, torch_q.QuantStub):
                    m.apply(torch.quantization.enable_fake_quant)
                    m.apply(torch.quantization.enable_observer)

                elif hasattr(m, 'activation_post_process'):
                    m.activation_post_process.apply(torch.quantization.disable_fake_quant)
                    m.activation_post_process.apply(torch.quantization.disable_observer)

                    if hasattr(m, 'weight_fake_quant'):
                        m.weight_fake_quant.apply(torch.quantization.enable_fake_quant)
                        m.weight_fake_quant.apply(torch.quantization.enable_observer)
                else:
                    for n, c in m.named_children():
                        q.put((f'{s}{n}', c))

        if self.backend == 'onnx':
            self.leaf_nodes = []
            self.swap_nodes = []

            q = queue.Queue()
            for node in graph.output_nodes:
                for prev_node in node.prev_nodes:
                    q.put((prev_node, False))

            while not q.empty():
                n, state = q.get()
                if isinstance(n.module, nn.Module):
                    orig_name = graph.module_original_name_dict.get(id(n.module))
                    new_mod, _ = graph.get_submodule_with_parent_from_name(orig_name, self.inplace)
                    if isinstance(new_mod, torch_q.DeQuantStub):
                        state = True
                    else:
                        if state:
                            if isinstance(new_mod, torch_q.QuantStub):
                                state = False
                            elif isinstance(new_mod, nn.Module) and hasattr(new_mod, 'activation_post_process'):
                                self.leaf_nodes.append(new_mod)
                                state = False
                            elif (
                                isinstance(new_mod, nn.Sequential)
                                and type(new_mod).__module__.startswith(nni.__name__)
                                and len(new_mod) > 0
                                and hasattr(new_mod[-1], 'activation_post_process')
                            ):
                                self.leaf_nodes.append(new_mod[-1])
                                state = False
                    for pn in n.prev_nodes:
                        q.put((pn, state))

            q = queue.Queue()
            visited = set()
            for node in graph.input_nodes:
                q.put((node, None, False, 0))

            while not q.empty():
                n, prev_q_mod, state, idx = q.get()
                key = f'{n.unique_name}:{idx}'
                if key in visited:
                    continue
                else:
                    visited.add(key)

                q_mod = prev_q_mod
                if n.quantized:
                    if isinstance(n.module, nn.Module):
                        orig_name = graph.module_original_name_dict.get(id(n.module))
                        new_mod, _ = graph.get_submodule_with_parent_from_name(orig_name, self.inplace)
                        if isinstance(new_mod, nn.Module) and hasattr(new_mod, 'activation_post_process'):
                            q_mod = new_mod
                        elif (
                            isinstance(new_mod, nn.Sequential)
                            and type(new_mod).__module__.startswith(nni.__name__)
                            and len(new_mod) > 0
                            and hasattr(new_mod[-1], 'activation_post_process')
                        ):
                            q_mod = new_mod[-1]
                        elif isinstance(new_mod, torch_q.DeQuantStub):
                            q_mod = new_mod
                        elif type(new_mod) is not nn.Identity:
                            state = True
                    else:
                        is_prev_float_functional = (
                            len(n.prev_nodes) > 1 and n.prev_nodes[0].type() is torch.nn.quantized.FloatFunctional
                        )
                        if is_prev_float_functional:
                            q_mod = getattr(n.prev_nodes[0].module, n.kind())
                        else:
                            state = True

                    if state and prev_q_mod is not None and q_mod != prev_q_mod:
                        self.swap_nodes.append((prev_q_mod, q_mod, idx))
                        state = False

                for next_n in n.next_nodes:
                    idx = next_n.prev_nodes.index(n)
                    q.put((next_n, q_mod, state, idx))

        if self.quantize_op_action.get(nn.LSTM, None) and self.backend == 'qnnpack':
            self.fake_quantize_lstm_weights()

        return graph.module

    def extra_qat_fusion_postprocess(self, graph):
        # Process additional fusable nodes
        processed_extra_q_rules = load_processed_extra_q_rules()
        is_extra_fusable = functools.partial(
            self.is_fusable,
            current_rules=processed_extra_q_rules,
            graph=graph,
        )

        custom_data = ([], set())
        graph.filter_forward_nodes(is_extra_fusable, custom_data, reverse=True)
        quant_list = custom_data[0]

        log.debug(f'Extra q postprocess for nodes: {quant_list}')

        rev_dict = dict((v, k) for k, v in graph.module_original_name_dict.items())
        for quant_nodes in quant_list:
            for orig_name in quant_nodes:
                if orig_name in rev_dict:
                    new_mod, _ = graph.get_submodule_with_parent_from_name(orig_name, self.inplace)
                    acp = getattr(new_mod, 'activation_post_process', None)
                    if acp is not None:
                        torch.quantization.disable_fake_quant(acp)
                        torch.quantization.disable_observer(acp)

                        activ_name = quant_nodes[-1]
                        if activ_name in rev_dict:
                            unique_name = graph.module_unique_name_dict[rev_dict[activ_name]]
                        else:
                            unique_name = activ_name
                            activ_name = None
                        node = graph.nodes_map[unique_name]
                        activ_type = node.kind()

                        post_dq = node.next_nodes[0].module
                        post_q = node.next_nodes[0].next_nodes[0].module

                        assert isinstance(post_dq, torch_q.DeQuantStub)
                        assert isinstance(post_q, torch_q.QuantStub)

                        dq_name = graph.module_original_name_dict[id(post_dq)]
                        q_name = graph.module_original_name_dict[id(post_q)]

                        post_acp = getattr(post_q, 'activation_post_process', None)
                        assert post_acp is not None

                        self.extra_qparams_mappings.append([acp, post_acp, dq_name, q_name, activ_name, activ_type])

    def disable_requantization_for_cat_pass(self, graph):
        def _find_quantized_cat_nodes(node: TraceNode, custom_node):
            # Find quantized cat nodes
            return node.type() == 'cat' and node.quantized

        # For cat nodes, the `activation_post_process` around it needs to be unified
        quantized_cat_nodes = graph.filter_forward_nodes(_find_quantized_cat_nodes)

        q = queue.Queue()
        visited_center = set()
        for n in quantized_cat_nodes:
            q.put((n, 'both', 0))
            parents = []
            names = []
            props = []
            visited_other = dict()
            while not q.empty():
                n, mode, fq_count = q.get()
                if (
                    n.kind() in ('shape', 'size')
                    or n.unique_name in visited_center
                    or visited_other.get(n.unique_name, 2) <= fq_count
                ):
                    continue

                if n.type() == 'cat':
                    visited_center.add(n.unique_name)
                else:
                    visited_other[n.unique_name] = fq_count

                new_fq_count = fq_count

                if isinstance(n.module, nn.Module):
                    is_prev_float_functional = False
                    orig_name = graph.module_original_name_dict.get(id(n.module))
                    new_mod, parent = graph.get_submodule_with_parent_from_name(orig_name, self.inplace)
                    prop = orig_name.split('.')[-1]
                    if QuantizableLSTM is not None and isinstance(new_mod, QuantizableLSTM):
                        if new_fq_count == 0:
                            if new_mod.bidirectional is False:
                                parents.append(new_mod.layers[-1].layer_fw.cell.ogate_cy)
                                names.append(f'{orig_name}.layer_fw.cell.ogate_cy.activation_post_process')
                                props.append('activation_post_process')
                            else:
                                parents.append(new_mod.layers[-1].layer_fw.cell.ogate_cy)
                                names.append(f'{orig_name}.layer_fw.cell.ogate_cy.activation_post_process')
                                props.append('activation_post_process')
                                parents.append(new_mod.layers[-1].layer_bw.cell.ogate_cy)
                                names.append(f'{orig_name}.layer_bw.cell.ogate_cy.activation_post_process')
                                props.append('activation_post_process')
                        new_fq_count += 1
                    elif QuantizableGRU is not None and isinstance(new_mod, QuantizableGRU):
                        if new_fq_count == 0:
                            if new_mod.bidirectional is False:
                                parents.append(new_mod.layers[-1].layer_fw.cell.add4)
                                names.append(f'{orig_name}.layer_fw.cell.add4.activation_post_process')
                                props.append('activation_post_process')
                            else:
                                parents.append(new_mod.layers[-1].layer_fw.cell.add4)
                                names.append(f'{orig_name}.layer_bw.cell.add4.activation_post_process')
                                props.append('activation_post_process')
                                parents.append(new_mod.layers[-1].layer_bw.cell.add4)
                                names.append(f'{orig_name}.layer_bw.cell.add4.activation_post_process')
                                props.append('activation_post_process')
                        new_fq_count += 1
                    elif isinstance(new_mod, (torch_q.FakeQuantize, torch_q.ObserverBase)):
                        if new_fq_count == 0:
                            parents.append(parent)
                            names.append(orig_name)
                            props.append(prop)
                        new_fq_count += 1
                    elif hasattr(new_mod, 'activation_post_process'):
                        if new_fq_count == 0:
                            parents.append(new_mod)
                            names.append(f'{orig_name}.activation_post_process')
                            props.append('activation_post_process')
                        new_fq_count += 1
                    elif (
                        isinstance(new_mod, nn.Sequential)
                        and type(new_mod).__module__.startswith(nni.__name__)
                        and len(new_mod) > 0
                        and hasattr(new_mod[-1], 'activation_post_process')
                    ):
                        if new_fq_count == 0:
                            parents.append(new_mod[-1])
                            names.append(f'{orig_name}[-1].activation_post_process')
                            props.append('activation_post_process')
                        new_fq_count += 1
                    if isinstance(new_mod, (torch_q.DeQuantStub, torch_q.QuantStub)):
                        new_fq_count = 2
                else:
                    is_prev_float_functional = (
                        len(n.prev_nodes) > 1 and n.prev_nodes[0].type() is torch.nn.quantized.FloatFunctional
                    )
                    if n.type() == 'cat':
                        mode = 'both'
                        fq_count = 0
                        new_fq_count = 0
                    if is_prev_float_functional:
                        m = n.prev_nodes[0].module
                        orig_name = graph.module_original_name_dict.get(id(m))
                        if new_fq_count == 0:
                            parents.append(m)
                            names.append(f'{orig_name}.activation_post_process')
                            props.append('activation_post_process')
                        new_fq_count += 1

                if mode in ('both', 'down'):
                    fq_up = fq_count
                    fq_down = new_fq_count
                elif mode == 'up':
                    fq_up = new_fq_count
                    fq_down = fq_count

                if mode == 'up' and len(n.next_nodes) > 1:
                    mode = 'both'
                    fq_down += 1

                if mode in ('both', 'up'):
                    for i, node in enumerate(n.prev_nodes):
                        if is_prev_float_functional and i == 0:
                            continue
                        if fq_up < 2:
                            q.put((node, 'up', fq_up))
                if mode in ('both', 'down'):
                    for node in n.next_nodes:
                        if fq_down < 2:
                            q.put((node, 'down', fq_down))

            if len(names) > 1:
                log.debug(f'Unifying the following nodes into one: {", ".join(names)}')
                unified = getattr(parents[0], props[0])
                for parent, prop in zip(parents[1:], props[1:]):
                    setattr(parent, prop, unified)

    def prepare_quantized_inputs_pass(self, graph):
        if self.quantized_input_stats is not None:
            assert len(self.quantized_input_stats) == len(graph.input_nodes), (
                f"quantized_input_stats contains {len(self.quantized_input_stats)} elements, but"
                f" {len(graph.input_nodes)} is expected"
            )

            for qstats, node in zip(self.quantized_input_stats, graph.input_nodes):
                if qstats is not None:
                    assert (
                        isinstance(qstats, (list, tuple))
                        and len(qstats) == 2
                        and all((isinstance(q, (int, float)) for q in qstats))
                    ), "quantized_input_stats format: [(mean_1, std_1), (mean_2, std_2), ..., (mean_n, std_n)]"

                for next_node in node.next_nodes:
                    if isinstance(next_node.module, torch_q.QuantStub):
                        self.set_module_quantization(next_node.module, qstats[0], qstats[1])

    def prepare_quantized_ops_pass(self, graph):
        if self.quantized_op_stats:

            def _find_quantized_ops_with_type(node, custom_data):
                if isinstance(node.module, nn.Module) and hasattr(node.module, 'activation_post_process'):
                    qstats = self.quantized_op_stats.get(node.kind(), None)
                    if qstats is not None:
                        custom_data.append((qstats, node))
                        return True
                    else:
                        if isinstance(node.module, torch_q.QuantStub):
                            qstats = self.quantized_op_stats.get(node.prev_nodes[0].kind(), None)
                            if qstats is not None:
                                custom_data.append((qstats, node))
                                return True
                return False

            node_with_qstats = []
            graph.filter_forward_nodes(_find_quantized_ops_with_type, node_with_qstats)

            for qstats, node in node_with_qstats:
                if qstats is not None:
                    assert (
                        isinstance(qstats, (list, tuple))
                        and len(qstats) == 2
                        and all((isinstance(q, (int, float)) for q in qstats))
                    ), (
                        "quantized_op_stats format: {'op_0': (mean_1, std_1), 'op_1': (mean_2, std_2), ..., 'op_n':"
                        " (mean_n, std_n)}"
                    )

                self.set_module_quantization(node.module, qstats[0], qstats[1])

    def set_module_quantization(self, module: nn.Module, mean: int, std: float):
        acp = getattr(module, 'activation_post_process', None)
        if acp is not None:
            fq_base_cls = getattr(torch_q, 'FakeQuantizeBase', torch_q.FakeQuantize)
            if isinstance(acp, fq_base_cls):
                fake_quant = acp

                module.apply(torch_q.disable_observer)

                scale = torch.tensor(1.0 / std, dtype=fake_quant.scale.dtype)
                offset = torch.tensor(mean, dtype=fake_quant.zero_point.dtype)

                fake_quant.scale.copy_(scale)
                fake_quant.zero_point.copy_(offset)

                quant_min = fake_quant.quant_min
                quant_max = fake_quant.quant_max

                observer = getattr(fake_quant, 'activation_post_process', None)
            elif isinstance(acp, torch_q.ObserverBase):
                observer = acp

                # We cannot use `disable_observer` which is designed for `FakeQuant` modules`.
                # Instead, we need to monkey-patch the forward function of the observer.
                identity = nn.Identity()
                observer.forward = identity.forward

                scale = torch.tensor(1.0 / std, dtype=torch.float32)
                offset = torch.tensor(mean, dtype=torch.int32)

                quant_min = observer.quant_min
                quant_max = observer.quant_max

            else:
                log.warning(f'Given module {type(module)} doesn\'t seem to be a quantized module')
                return

            if observer is not None:
                if quant_min is None and quant_max is None:
                    if observer.reduce_range:
                        quant_min, quant_max = 0, 127
                    else:
                        quant_min, quant_max = 0, 255

                observer.min_val = scale * (quant_min - offset)
                observer.max_val = scale * (quant_max - offset)

    def rescale_activations_with_quant_min_max(self, quant_min: int, quant_max: int) -> None:
        """Rescales activations with provided quant_min and quant_max"""
        for n, m in self.model.named_modules():
            if '.weight_fake_quant' in n:
                continue

            if isinstance(m, torch.nn.Sigmoid) and quant_min == 0 and quant_max == 65535:
                m.activation_post_process.scale = m.activation_post_process.scale * 2
                m.activation_post_process.zero_point = torch.zeros_like(m.activation_post_process.zero_point) + 32768
                log.info(f"Rescale {n}(sigmoid mod) quant param to [1/32768, 0] to align TFLite.")

            if isinstance(m, torch.quantization.FakeQuantize):
                observer = getattr(m, 'activation_post_process', None)
                if observer is not None:
                    old_quant_range = m.quant_max - m.quant_min + 1
                    m.quant_min = quant_min
                    m.quant_max = quant_max
                    new_quant_range = m.quant_max - m.quant_min + 1
                    m.scale = m.scale * old_quant_range / new_quant_range
                    zero_point = new_quant_range // 2 if m.zero_point != 0 else 0
                    m.zero_point = torch.zeros_like(m.zero_point) + zero_point
                    observer.quant_min = quant_min
                    observer.quant_max = quant_max

    def rewrite_quantize_graph_for_tensorrt(self, graph: TraceGraph) -> None:
        """Rewrites the computation graph for TensorRT quantization"""

        processed_types = set()
        for types in TENSORRT_OBSERVED_NODES:
            if isinstance(types, tuple):
                if hasattr(types[0], types[1]):
                    processed_types.add(getattr(types[0], types[1]))
            else:
                processed_types.add(types)

        def _is_observed_nodes_for_tensorrt(node, custom_data):
            return node.kind() in processed_types

        observed_nodes = graph.filter_forward_nodes(_is_observed_nodes_for_tensorrt)
        node_map = dict()
        for idx, node in enumerate(observed_nodes):
            fake_quant_cls = torch_q.QuantStub
            assert node.rev_index is False
            prev_nodes = {n.unique_name: n for n in node.prev_nodes}.values()
            for inner_idx, prev_node in enumerate(prev_nodes):

                if prev_node.kind() in ('shape', 'device', 'size', 'dtype'):
                    continue

                prev_tensor_ptrs = []
                for pt in node.prev_tensors:
                    for nt in prev_node.next_tensors:
                        if isinstance(nt, (list, tuple)):
                            for k, ntt in enumerate(nt):
                                if id(pt) == id(ntt):
                                    if id(pt) not in prev_tensor_ptrs:
                                        prev_tensor_ptrs.append(id(pt))
                                    break
                        elif id(pt) == id(nt):
                            if id(pt) not in prev_tensor_ptrs:
                                prev_tensor_ptrs.append(id(pt))
                            break

                for ptr_idx, ptr in enumerate(prev_tensor_ptrs):
                    if ptr in node_map:
                        fake_quant = node_map[ptr]

                        graph.insert_between(prev_node, node, fake_quant, move_idx=True, tensor_ptrs=set([ptr]))
                    else:

                        fake_quant = fake_quant_cls()

                        fake_quant_name = f'fake_quant_inner_{idx}_{inner_idx}_{ptr_idx}'

                        graph.module_unique_name_dict[id(fake_quant)] = fake_quant_name
                        graph.module_original_name_dict[id(fake_quant)] = fake_quant_name

                        module_constructor_lines[id(fake_quant)] = f'{qualified_name(fake_quant_cls)}()'

                        graph.insert_between(prev_node, node, fake_quant, move_idx=True, tensor_ptrs=set([ptr]))
                        node_map[ptr] = graph.nodes_map[fake_quant_name]

        graph.quantized = True
        graph.recompute_forward_order()

    def rewrite_quantize_graph(self, graph: TraceGraph) -> None:
        """Rewrites the computation graph for quantization"""
        if graph.quantized:
            return

        graph_quantized = False
        for n in graph.forward_nodes:
            if n.type() in (torch_q.QuantStub, torch_q.DeQuantStub):
                graph_quantized = True
                break

        for n in graph.other_init_nodes:
            if n.type() is nnq.FloatFunctional:
                graph_quantized = True
                break

        if graph_quantized:
            log.warning(
                'Graph is quantized. No need to rewrite. Please pass in `config={"rewrite_graph": False}` to suppress'
                ' this warning'
            )
            return

        if self.backend == 'tensorrt':
            return self.rewrite_quantize_graph_for_tensorrt(graph)

        if self.layerwise_default is False:
            for n in graph.forward_nodes:
                self.layerwise_config.setdefault(n.unique_name, False)

        creation_func_names = load_creation_func_names()

        def _is_extra_constant_nodes(node, custom_data):
            return node.full_name() in creation_func_names

        extra_constant_nodes = graph.filter_forward_nodes(_is_extra_constant_nodes)

        def _is_int_to_float_nodes(node, custom_data):
            if node.full_name() in creation_func_names:
                return False

            if len(node.prev_nodes) == 1 and len(node.next_nodes) == 1:
                if len(node.prev_tensors) == 1 and len(node.next_tensors) == 1:
                    if node.prev_nodes[0].kind() == 'shape' and node.prev_nodes[0].module.is_property:
                        return False
                    if (
                        isinstance(node.prev_tensors[0], torch.Tensor)
                        and isinstance(node.next_tensors[0], torch.Tensor)
                        and node.prev_tensors[0].dtype in (torch.int32, torch.int64)
                        and torch.is_floating_point(node.next_tensors[0])
                    ):
                        return True
            else:
                return False

        int_to_float_nodes = graph.filter_forward_nodes(_is_int_to_float_nodes)

        # When converting float-tensor to int16/int32/int64, we need to add 'fake_dequant' node before convert-node.
        def _is_float_to_non_float_nodes(node, custom_data):
            if isinstance(node.module, TraceFunction) and node.module.kind in ('shape', 'size'):
                return False
            return (
                len(node.next_tensors) == 1
                and len(node.prev_tensors) > 0
                and isinstance(node.prev_tensors[0], torch.Tensor)
                and isinstance(node.next_tensors[0], torch.Tensor)
                and torch.is_floating_point(node.prev_tensors[0])
                and not torch.is_floating_point(node.next_tensors[0])
            )

        float_to_non_float_nodes = graph.filter_forward_nodes(_is_float_to_non_float_nodes)

        def _is_params_in_module(node, custom_data):
            if len(node.prev_nodes) == 1 and len(node.next_nodes) == 1:
                if len(node.prev_tensors) == 1 and len(node.next_tensors) == 1:
                    if isinstance(node.prev_tensors[0], nn.Module) and not isinstance(
                        node.prev_tensors[0], nnq.FloatFunctional
                    ):
                        return True
            return False

        param_nodes = graph.filter_forward_nodes(_is_params_in_module)
        for node in param_nodes:
            prev_node = node.prev_nodes[0]
            is_known_mod = prev_node.kind().__name__ in (
                'Conv1d',
                'Conv2d',
                'Linear',
                'ConvTranspose1d',
                'ConvTranspose2d',
            )
            prop = node.module.full_name

            if is_known_mod and prop in ('weight', 'bias') and self.layerwise_config.get(prev_node.unique_name, True):
                node.module.is_class = False
                node.module.is_property = False

                node.module.full_name = 'tinynn.graph.quantization.utils.get_parameter'
                node.module.args_template = f'{{}}, "{prop}"'
                node.module.args_template_no_self = f'"{prop}"'
                node.module.args_offset = 0
                node.module.update_args_string()

        # First, we insert the QuantStub nodes for every input/constant node
        for idx, node in reversed(
            list(
                enumerate(
                    graph.input_nodes + graph.constant_nodes + extra_constant_nodes + int_to_float_nodes + param_nodes
                )
            )
        ):
            if node.next_tensors[0].dtype in (torch.int32, torch.int64):
                continue

            fake_quant = torch_q.QuantStub()

            graph.module_unique_name_dict[id(fake_quant)] = f'fake_quant_{idx}'
            graph.module_original_name_dict[id(fake_quant)] = f'fake_quant_{idx}'

            fake_quant_cls = type(fake_quant)
            module_constructor_lines[id(fake_quant)] = f'{qualified_name(fake_quant_cls)}()'

            graph.insert_after(node, fake_quant)

        # Second, we insert the DeQuantStub nodes for every output node
        for idx, node in enumerate(graph.output_nodes + float_to_non_float_nodes):
            fake_dequant_cls = torch_q.DeQuantStub
            if node.rev_index:
                modules = []
                for rev_idx in range(len(node.prev_nodes)):
                    if node.prev_tensors[rev_idx].dtype in (torch.int32, torch.int64):
                        fake_dequant = nn.Identity()
                    else:
                        fake_dequant = fake_dequant_cls()

                    graph.module_unique_name_dict[id(fake_dequant)] = f'fake_dequant_{idx}_{rev_idx}'
                    graph.module_original_name_dict[id(fake_dequant)] = f'fake_dequant_{idx}_{rev_idx}'

                    module_constructor_lines[id(fake_dequant)] = f'{qualified_name(fake_dequant_cls)}()'
                    modules.append(fake_dequant)

                graph.insert_before(node, modules, move_idx=True)
            else:
                if node.prev_tensors[0].dtype in (torch.int32, torch.int64):
                    continue

                fake_dequant = fake_dequant_cls()

                graph.module_unique_name_dict[id(fake_dequant)] = f'fake_dequant_{idx}'
                graph.module_original_name_dict[id(fake_dequant)] = f'fake_dequant_{idx}'

                module_constructor_lines[id(fake_dequant)] = f'{qualified_name(fake_dequant_cls)}()'

                if len(node.prev_nodes) > 1:
                    # Insert 'fake_dequant' node before type conversion operators
                    graph.insert_between(node.prev_nodes[0], node, fake_dequant, move_idx=True)
                else:
                    graph.insert_before(node, fake_dequant, move_idx=True)

        # Third, we rewrite neg/sub/div using supported functions(e.g add, mul)
        def _is_neg_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return (
                    cur_module.kind == 'neg'
                    and torch.is_floating_point(cur_module.prev_tensors[0])
                    and self.layerwise_config.get(node.unique_name, True)
                )

        neg_nodes = graph.filter_forward_nodes(_is_neg_node)
        log.info(f'rewriting neg for {[node.unique_name for node in neg_nodes]}')
        for idx, node in enumerate(neg_nodes):
            node.module.func_type = '__mul__'
            node.module.kind = 'mul'

            full_name_parts = node.module.full_name.split('.')
            full_name_parts[-1] = node.module.func_type

            node.module.full_name = '.'.join(full_name_parts)

            with override_current_trace_graph(graph):
                node.module.parse_args(node.prev_tensors[0], -1.0)

        def _is_div_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            # Current, only the following condition could be handled.
            #   a / constant => a * (1 / constant)
            if cur_class == TraceFunction:
                return (
                    (cur_module.kind == 'truediv' or cur_module.func_type in ('div', 'div_'))
                    and len(cur_module.prev_tensors) == 1
                    and torch.is_floating_point(cur_module.prev_tensors[0])
                    and cur_module.func_type != '__rtruediv__'
                    and torch.is_floating_point(node.next_tensors[0])
                    and node.prev_nodes[0].kind() not in ('size', 'shape')
                    and self.layerwise_config.get(node.unique_name, True)
                )

        div_nodes = graph.filter_forward_nodes(_is_div_node)
        log.info(f'rewriting div for {[node.unique_name for node in div_nodes]}')
        for idx, node in enumerate(div_nodes):
            op_type = node.module.func_type
            inplace = op_type in ('__itruediv__', 'itruediv', 'div_')

            if inplace:
                node.module.func_type = '__imul__'
            else:
                node.module.func_type = '__mul__'

            node.module.kind = 'mul'

            full_name_parts = node.module.full_name.split('.')
            full_name_parts[-1] = node.module.func_type

            node.module.full_name = '.'.join(full_name_parts)

            # Here we make a simple guess, if dot is in the string, then it's a floating number.
            # Otherwise, it is an integral number.
            if '.' in node.module.args_string_no_self:
                other_arg = float(node.module.args_string_no_self)
            else:
                other_arg = int(node.module.args_string_no_self)

            with override_current_trace_graph(graph):
                node.module.parse_args(node.prev_tensors[0], 1.0 / other_arg)

        def _is_sub_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return (
                    cur_module.kind == 'sub'
                    and torch.is_floating_point(cur_module.prev_tensors[0])
                    and torch.is_floating_point(node.next_tensors[0])
                    and self.layerwise_config.get(node.unique_name, True)
                )

        sub_nodes = graph.filter_forward_nodes(_is_sub_node)
        log.info(f'rewriting sub for {[node.unique_name for node in sub_nodes]}')
        for idx, node in enumerate(sub_nodes):
            op_type = node.module.func_type

            full_name_parts = node.module.full_name.split('.')
            full_name_parts[-1] = node.module.func_type

            if len(node.module.prev_tensors) == 1 and node.module.func_type != '__rsub__':
                node.module.func_type = '__add__'
                node.module.kind = 'add'
                node.module.full_name = '.'.join(full_name_parts)

                # Here we make a simple guess, if dot is in the string, then it's a floating number.
                # Otherwise, it is an integral number.
                if '.' in node.module.args_string_no_self or 'e' in node.module.args_string_no_self:
                    other_arg = float(node.module.args_string_no_self)
                else:
                    other_arg = int(node.module.args_string_no_self)

                with override_current_trace_graph(graph):
                    node.module.parse_args(node.prev_tensors[0], -other_arg)
            elif len(node.module.prev_tensors) == 2 and len(node.prev_nodes) == 2:
                new_fullname_parts = copy.deepcopy(full_name_parts)
                new_fullname_parts[-1] = '__mul__'
                new_fullname = '.'.join(new_fullname_parts)
                current_tensor = node.prev_tensors[0]
                input_node = node.prev_nodes[1]
                input_tensor = node.prev_tensors[1]
                output_tensor = input_tensor * -1

                with override_current_trace_graph(graph):
                    trace_func = TraceFunction(new_fullname, True, prefix='fuse_').parse_args(input_tensor, -1)
                    graph.insert_between(input_node, node, trace_func, [output_tensor], True)

                    node.module.func_type = '__add__'
                    node.module.kind = 'add'

                    full_name_parts[-1] = node.module.func_type

                    node.module.full_name = '.'.join(full_name_parts)

                    node.module.parse_args(current_tensor, output_tensor)

            elif node.module.func_type == '__rsub__' and len(node.module.prev_tensors) == 1:
                new_fullname_parts = copy.deepcopy(full_name_parts)
                new_fullname_parts[-1] = '__mul__'
                new_fullname = '.'.join(new_fullname_parts)
                input_node = node.prev_nodes[0]
                input_tensor = node.prev_tensors[0]
                output_tensor = input_tensor * -1

                if '.' in node.module.args_string_no_self:
                    other_arg = float(node.module.args_string_no_self)
                else:
                    other_arg = int(node.module.args_string_no_self)

                with override_current_trace_graph(graph):
                    trace_func = TraceFunction(new_fullname, True, prefix='fuse_').parse_args(input_tensor, -1)
                    graph.insert_between(input_node, node, trace_func, [output_tensor], True)

                    node.module.func_type = '__radd__'
                    node.module.kind = 'add'

                    full_name_parts[-1] = node.module.func_type

                    node.module.full_name = '.'.join(full_name_parts)

        # Then, we write torch.stack nodes to torch.unsqueeze + torch.cat
        def _is_stack_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            return (
                cur_class == TraceFunction
                and cur_module.kind == 'stack'
                and torch.is_floating_point(node.next_tensors[0])
                and self.layerwise_config.get(node.unique_name, True)
            )

        stack_nodes = graph.filter_forward_nodes(_is_stack_node)

        for idx, node in enumerate(stack_nodes):
            args = getattr(node.module, 'args_string_no_self', '')
            if ',' in args:
                log.error('rewrite doesn\'t support multiple args for torch.stack')
                assert False
            if len(args) > 0:
                if '=' in args:
                    k, v = args.split('=')
                    assert k in ('axis', 'dim')
                    dim = int(v)
                    if k == 'axis':
                        node.module.args_template_no_self = node.module.args_template_no_self.replace('axis', 'dim')
                        node.module.args_template = node.module.args_template.replace('axis', 'dim')
                        node.module.update_args_string()
                else:
                    dim = int(args)
            else:
                dim = 0
            unique_prev_nodes = {n.unique_name: n for n in node.prev_nodes}.values()
            for n in unique_prev_nodes:
                shared_tensors = list(set(node.prev_tensors).intersection(set(n.next_tensors)))
                if len(shared_tensors) == 0:
                    log.debug('tensor rewrite already done, skipping')
                    continue
                for t in shared_tensors:
                    with override_current_trace_graph(graph):
                        trace_func = TraceFunction('torch.unsqueeze', prefix='fuse_').parse_args(t, dim)
                    next_tensors = [torch.unsqueeze(t, dim)]
                    graph.insert_between(n, node, trace_func, next_tensors, move_idx=True, tensor_ptrs=set([id(t)]))

            node.module.func_type = 'cat'
            node.module.kind = 'cat'

            full_name_parts = node.module.full_name.split('.')
            full_name_parts[-1] = node.module.func_type

            node.module.full_name = '.'.join(full_name_parts)

        # Next, we rewrite add/mul/cat with one float32 output using torch.nn.quantized.FloatFunctional
        def _is_convertible_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return (
                    cur_module.kind in ('add', 'mul', 'cat')
                    and torch.is_floating_point(node.next_tensors[0])
                    and self.layerwise_config.get(node.unique_name, True)
                    and all((torch.is_floating_point(pt) for pt in node.prev_tensors))
                )

        convertible_nodes = graph.filter_forward_nodes(_is_convertible_node)
        unfusable_add_nodes = []  # noqa: F841
        log.info(f'rewriting add/mul/cat for {[node.unique_name for node in convertible_nodes]}')
        for idx, node in enumerate(convertible_nodes):
            old_full_name = node.module.full_name  # noqa: F841
            old_is_class = node.module.is_class  # noqa: F841

            op_kind = node.module.kind
            float_functional = nnq.FloatFunctional()

            module_name = f'float_functional_simple_{idx}'

            graph.module_unique_name_dict[id(float_functional)] = module_name
            graph.module_original_name_dict[id(float_functional)] = module_name

            float_functional_cls = type(float_functional)
            module_constructor_lines[id(float_functional)] = f'{qualified_name(float_functional_cls, short=True)}()'

            new_node = TraceNode(float_functional, cur_graph=graph)
            graph.nodes_map[new_node.unique_name] = new_node
            graph.other_init_nodes.append(new_node)

            node.module.is_class = False
            prev_tensor_size = len(node.prev_tensors)
            if op_kind in ('add', 'mul'):
                if prev_tensor_size == 2:
                    if node.prev_nodes[1].prev_nodes[0].kind() not in ('shape', 'size'):
                        op_type = op_kind
                    else:
                        op_type = f'{op_kind}_scalar'
                elif prev_tensor_size == 1:
                    op_type = f'{op_kind}_scalar'
                else:
                    log.error(f'Unknown add/mul type for {node.unique_name}, prev tensor size: {prev_tensor_size}')
                    assert False
            else:
                # Don't check anything for other OPs.
                # It is simply too complex for us.
                op_type = op_kind
            node.module.func_type = op_type
            node.module.full_name = f'self.{module_name}.{op_type}'
            # Inplace operations
            if node.module.func_type in ['__iadd__', '__imul__', 'add_', 'mul_']:
                node.module.add_alias(node.module.tensor_names[0])
                q = queue.Queue()
                q.put(node.prev_nodes[0])
                while not q.empty():
                    n = q.get()
                    if type(n.module) is TraceFunction:
                        prev_aliases = n.module.get_aliases()
                        if prev_aliases is not None:
                            for pa in reversed(prev_aliases):
                                node.module.add_alias(pa, head=True)
                    else:
                        if getattr(n.module, 'inplace', False):
                            q.put(n.prev_nodes[0])
                            node.module.add_alias(n.prev_node_unique_name(0), head=True)

            # We need to convert radd to normal add here
            if node.module.func_type in ['__radd__', '__rmul__']:
                if '=' in node.module.args_string_no_self or ', ' in node.module.args_string_no_self:
                    log.error(f'Don\'t know how to translate {node.module.args_string_no_self} for __radd__/__rmul__')
                    assert False
                if prev_tensor_size == 1:
                    # Here we make a simple guess, if dot is in the string, then it's a floating number.
                    # Otherwise, it is an integral number.
                    if '.' in node.module.args_string_no_self or 'e' in node.module.args_string_no_self:
                        other_arg = float(node.module.args_string_no_self)
                    else:
                        other_arg = int(node.module.args_string_no_self)

                    with override_current_trace_graph(graph):
                        node.module.parse_args(node.prev_tensors[0], other_arg)
                else:
                    with override_current_trace_graph(graph):
                        # It is even simple here. We only need to swap the order of the tensors.
                        node.module.parse_args(node.prev_tensors[1], node.prev_tensors[0])

        # Rewrite torch.clamp_{min,max} to torch.clamp
        def _is_clamp_min_max_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return (
                    cur_module.kind in ('clamp_min', 'clamp_max')
                    and torch.is_floating_point(node.next_tensors[0])
                    and self.layerwise_config.get(node.unique_name, True)
                )

        clamp_min_max_nodes = graph.filter_forward_nodes(_is_clamp_min_max_node)
        log.info(f'rewriting clamp_{{min,max}} for {[node.unique_name for node in convertible_nodes]}')
        for node in clamp_min_max_nodes:
            kw = node.module.kind[-3:]
            _parse_args = eval(f'lambda {kw}, *args, **kwargs: {kw}')

            val = eval(f'_parse_args({node.module.args_string_no_self})')

            if kw != 'min' or val != 0.0:
                continue

            if kw == 'min':
                arg_str = f'{val}'
            else:
                arg_str = f'None, {val}'

            sub_str = f'_{kw}'
            node.module.kind = node.module.kind.replace(sub_str, '')
            node.module.func_type = node.module.func_type.replace(sub_str, '')
            node.module.full_name = '.'.join(node.module.full_name.split('.')[:-1] + [node.module.func_type])

            node.module.args_template_no_self = arg_str
            node.module.args_template = f'{{}}, {arg_str}'
            node.module.update_args_string()

        # Rewrite torch.clamp to relu, relu6 and clamp_{min, max}
        def _is_clamp_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return (
                    cur_module.kind == 'clamp'
                    and torch.is_floating_point(node.next_tensors[0])
                    and self.layerwise_config.get(node.unique_name, True)
                )

        clamp_nodes = graph.filter_forward_nodes(_is_clamp_node)
        log.info(f'rewriting clamp for {[node.unique_name for node in convertible_nodes]}')
        for node in clamp_nodes:

            def _parse_args(min=None, max=None, *args, **kwargs):
                return min, max

            min, max = eval(f'_parse_args({node.module.args_string_no_self})')

            kind = None
            if min == 0.0:
                if max is None:
                    kind = 'relu'
                elif max == 6.0:
                    kind = 'relu6'

            if kind is None:
                if max is None:
                    kind = 'clamp_min'
                elif min is None:
                    kind = 'clamp_max'
                else:
                    kind = 'clamp_with_fusion'

            if kind in ('clamp_min', 'clamp_max'):
                node.module.kind = kind
                node.module.func_type = node.module.func_type.replace('clamp', kind)
                node.module.full_name = '.'.join(node.module.full_name.split('.')[:-1] + [node.module.func_type])

                if max is None:
                    arg_str = f'{min}'
                else:
                    arg_str = f'{max}'
            elif kind == 'clamp_with_fusion':
                node.module.is_class = False
                node.module.kind = kind
                node.module.func_type = node.module.func_type.replace('clamp', kind)
                node.module.full_name = f'tinynn.graph.quantization.utils.{node.module.func_type}'

                arg_str = f'{min}, {max}'
            else:
                inplace = node.module.func_type == f'{node.module.kind}_'
                node.module.kind = kind
                node.module.func_type = kind
                node.module.full_name = f'torch.nn.functional.{kind}'

                if inplace:
                    arg_str = 'inplace=True'
                else:
                    arg_str = ''

            node.module.args_template_no_self = arg_str
            node.module.args_template = f'{{}}, {arg_str}'
            node.module.update_args_string()

        # Rewrite other fusable functions
        # e.g. add_relu(x, y) =>
        #            r = torch.add(x, y)
        #            r = torch.nn.functional.relu(r)
        def _is_add_relu_fusable_node(node: TraceNode, custom_data) -> bool:
            cur_module = node.module
            cur_class = type(cur_module)
            visited_nodes = [node]

            if self.layerwise_config.get(node.unique_name, True) is False:
                return False

            if cur_class == TraceFunction:
                # The intermediate result cannot be used if fused.
                # So we don't fuse the nodes under such circumstances.
                if cur_module.kind != 'add' or len(node.next_nodes) != 1:
                    return False
                # The input for the add operations should be two tensors.
                if len(node.prev_tensors) != 2:
                    return False
                # We accept inplace operations for both add and relu.
                # The inplace property could be elinimated because we track the tensors
                # instead of their names.
                next_node = node.next_nodes[0]
                next_module = next_node.module
                next_class = type(next_module)
                if next_class == TraceFunction:
                    fuse = next_module.kind == 'relu'
                else:
                    while next_class == nn.Identity:
                        cur_node = next_node
                        visited_nodes.append(cur_node)
                        if len(cur_node.next_nodes) != 1:
                            return False
                        next_node = cur_node.next_nodes[0]
                        next_module = next_node.module
                        next_class = type(next_module)

                    fuse = next_class.__name__ == 'ReLU'

                if not fuse:
                    return False

                if type(next_node.module) is TraceFunction:
                    inplace = next_node.module.func_type == 'relu_' or 'True' in next_node.module.args_string
                else:
                    inplace = getattr(next_node.module, 'inplace', False)

                # Inplace check
                # If add is inplace and relu is not inplace, we need to ensure that all the aliases of
                # the first operand of add are not used when relu is called.
                if not inplace:
                    aliases = cur_module.get_aliases()
                    if aliases:
                        q = queue.Queue()
                        q.put(node.prev_nodes[0])
                        while not q.empty():
                            n = q.get()
                            last_order = max((x.forward_order for x in n.next_nodes))
                            if last_order > node.forward_order:
                                fuse = False
                                break
                            if type(n.module) is TraceFunction and n.module.get_aliases():
                                q.put(n.prev_nodes[0])
                            elif getattr(n.module, 'inplace', False):
                                q.put(n.prev_nodes[0])

                return fuse

        add_relu_fusable_nodes = graph.filter_forward_nodes(_is_add_relu_fusable_node)
        for node in add_relu_fusable_nodes:
            full_name = node.module.full_name.replace('add', 'add_relu')
            next_node = node.next_nodes[0]
            kind = 'add_relu'
            func_type = kind
            is_class = False
            nodes_to_fuse = [node, next_node]
            while next_node.type() is nn.Identity:
                next_node = next_node.next_nodes[0]
                nodes_to_fuse.append(next_node)
            if type(next_node.module) is TraceFunction:
                inplace = next_node.module.func_type == 'relu_' or 'True' in next_node.module.args_string
            else:
                inplace = next_node.module.inplace
            graph.fuse_nodes_to_func(nodes_to_fuse, full_name, kind, func_type, is_class)
            # Propagate aliases for inplace nodes
            if inplace:
                aliases = node.module.get_aliases()
                if aliases:
                    node.module.add_alias(node.module.tensor_names[0])
            else:
                node.module.aliases = None

        # Rewrite relu, relu6 as nn.ReLU() and nn.ReLU6() for Module fusable rules
        def _is_functional_rewrite_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return cur_module.kind in FUNCTIONAL_MODULE_MAPPING and self.layerwise_config.get(
                    node.unique_name, True
                )

        func_nodes_to_rewrite = graph.filter_forward_nodes(_is_functional_rewrite_node)
        log.info(f'rewriting functional to module for {[node.unique_name for node in func_nodes_to_rewrite]}')
        for idx, node in enumerate(func_nodes_to_rewrite):
            kind = node.module.kind
            inplace = node.module.func_type == f'{kind}_' or 'True' in node.module.args_string
            klass = FUNCTIONAL_MODULE_MAPPING[kind]
            arguments = getattr(klass, '__constants__', None)
            if arguments is None:
                new_func = klass()
            elif node.module.kind in ('relu', 'relu6', 'silu', 'hardswish'):
                new_func = klass(inplace=inplace)
            elif node.module.kind in ('elu', 'leaky_relu'):
                if hasattr(node.module, 'args_string_no_self'):

                    def _parse_args(alpha=1.0, *args, **kwargs):  # noqa: F811
                        return alpha

                    alpha = eval(f'_parse_args({node.module.args_string_no_self})')
                    if node.module.kind == 'leaky_relu':
                        new_func = klass(alpha, inplace=inplace)
                    else:
                        new_func = klass(alpha, inplace=inplace)
                else:
                    alpha = None
                    if node.module.kind == 'leaky_relu':
                        new_func = klass(inplace=inplace)
                    else:
                        new_func = klass()
            elif node.module.kind == 'prelu':
                weight_t = node.prev_tensors[1]
                weight_node = node.prev_nodes[1]
                if weight_node.type() is torch_q.QuantStub:
                    weight_node = weight_node.prev_nodes[0]

                if weight_node.type() is not ConstantNode or not weight_node.module.is_parameter:
                    log.warning('Rewrite for F.prelu(x, buffer) to nn.PReLU is skipped as it changes the semantics')
                    continue

                num_parameters = weight_t.nelement()
                new_func = klass(num_parameters)
                new_func.weight.data.copy_(node.prev_tensors[1])

                # Drop last input
                node.prev_tensors.pop()
                last_node = node.prev_nodes.pop()
                last_node.next_nodes.remove(node)

                # Remove unused forward functions iteratively
                new_last_node = last_node
                while new_last_node is not None:
                    last_node = new_last_node
                    if (
                        len(last_node.next_nodes) == 0
                        and len(last_node.prev_nodes) < 2
                        and last_node in graph.forward_nodes
                    ):
                        if len(last_node.prev_nodes) == 1:
                            new_last_node = last_node.prev_nodes[0]
                        else:
                            new_last_node = None
                        graph.remove_node(last_node)
                    else:
                        new_last_node = None
            elif node.module.kind == 'glu':
                if hasattr(node.module, 'args_string_no_self'):

                    def _parse_args_dim(dim=-1, *args, **kwargs):  # noqa: F811
                        return dim

                    dim = eval(f'_parse_args_dim({node.module.args_string_no_self})')
                    new_func = klass(dim)
                else:
                    dim = None
                    new_func = klass()
            else:
                raise NotImplementedError(f"Don't know how to parse {klass.__name__} with argument {arguments}")

            graph.module_unique_name_dict[id(new_func)] = f'rewritten_{kind}_{idx}'
            graph.module_original_name_dict[id(new_func)] = f'rewritten_{kind}_{idx}'

            relu_cls = type(new_func)
            if inplace:
                arg_str = 'inplace=True'
            else:
                arg_str = ''

            if node.module.kind in ('elu', 'leaky_relu') and alpha is not None:
                if arg_str:
                    arg_str = f'{alpha}, {arg_str}'
                else:
                    arg_str = f'{alpha}'

            if node.module.kind == 'prelu':
                if num_parameters != 1:
                    arg_str = f'{num_parameters}'
                else:
                    arg_str = ''

            if node.module.kind == 'glu':
                if dim is not None and dim != -1:
                    arg_str = f'{dim}'
                else:
                    arg_str = ''

            module_constructor_lines[id(new_func)] = f'{qualified_name(relu_cls)}({arg_str})'
            graph.replace_node_module(node, new_func)

        # Rewrite dropout as nn.Dropout() for models in training mode
        def _is_dropout_functional_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return cur_module.kind == 'dropout' and self.layerwise_config.get(node.unique_name, True)

        def _dropout_args(p=0.5, training=True, inplace=False):
            return p, training, inplace

        dropout_nodes_to_rewrite = graph.filter_forward_nodes(_is_dropout_functional_node)
        log.info(f'rewriting dropout for {[node.unique_name for node in dropout_nodes_to_rewrite]}')
        for idx, node in enumerate(dropout_nodes_to_rewrite):
            args = getattr(node.module, 'args_string_no_self', '')
            p, _, inplace = eval(f'_dropout_args({args})')
            kind = node.module.kind
            inplace = node.module.func_type == f'{kind}_' or inplace
            dropout_cls = nn.Dropout
            new_dropout = dropout_cls(p, inplace=inplace)

            graph.module_unique_name_dict[id(new_dropout)] = f'rewritten_{kind}_{idx}'
            graph.module_original_name_dict[id(new_dropout)] = f'rewritten_{kind}_{idx}'

            if inplace:
                arg_str = f'{p}, inplace={inplace}'
            else:
                arg_str = f'{p}'

            module_constructor_lines[id(new_dropout)] = f'{qualified_name(dropout_cls)}({arg_str})'
            graph.replace_node_module(node, new_dropout)

        # Add contiguous nodes for partially-supported OPs
        # Some of the operations support quantization, but they only accept contiguous input tensors.
        def _is_partially_quantizable(node, custom_data):
            if self.layerwise_config.get(node.unique_name, True) is False:
                return False
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == ConstantNode:
                return False
            elif cur_class == TraceFunction:
                return cur_module.kind in ('pad',)
            else:
                return cur_class in (nn.ConstantPad1d, nn.ConstantPad2d, nn.ConstantPad3d, nn.ZeroPad2d)

        if LooseVersion(torch.__version__) >= LooseVersion('1.7.0'):
            partially_supported_nodes = graph.filter_forward_nodes(_is_partially_quantizable)
            for idx, node in enumerate(partially_supported_nodes):
                for n in node.prev_nodes[:1]:
                    shared_tensors = list(set(node.prev_tensors).intersection(set(n.next_tensors)))
                    if len(shared_tensors) > 1:
                        log.error('rewrite for partially-supported ops supports with nodes with exact one input')
                        assert False
                    with override_current_trace_graph(graph):
                        trace_func = TraceFunction('torch.Tensor.contiguous', True, prefix='fuse_').parse_args(
                            shared_tensors[0]
                        )
                    next_tensors = [x.contiguous() for x in shared_tensors]
                    graph.insert_between(n, node, trace_func, next_tensors, True)

        # Remove non-leaf `.data` nodes
        def _is_non_leaf_data_nodes(node, custom_data):
            if self.layerwise_config.get(node.unique_name, True) is False:
                return False
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return cur_module.kind == 'data' and cur_module.is_property and len(node.next_nodes) > 0
            return False

        non_leaf_data_nodes = graph.filter_forward_nodes(_is_non_leaf_data_nodes)
        for idx, node in enumerate(non_leaf_data_nodes):
            graph.remove_node(node)

        # Handle PoolNd with kernel_size=1
        def _is_pool_nd_with_one_kernel_size(node, custom_data):
            if self.layerwise_config.get(node.unique_name, True) is False:
                return False
            cur_module = node.module
            cur_class = type(cur_module)
            kernel_size, stride = None, None
            if cur_class == TraceFunction:
                if cur_module.kind in ('avg_pool1d', 'avg_pool2d', 'max_pool1d', 'max_pool2d'):

                    def _avgpool_kernel_size_and_stride(kernel_size, stride=None, *args, **kwargs):
                        return kernel_size, stride

                    kernel_size, stride = eval(f'_avgpool_kernel_size_and_stride({cur_module.args_string_no_self})')
            else:
                if cur_class in (nn.AvgPool1d, nn.AvgPool2d, nn.MaxPool1d, nn.MaxPool2d):
                    kernel_size = cur_module.kernel_size
                    stride = cur_module.stride

            if kernel_size is not None:
                if isinstance(kernel_size, (tuple, list)):
                    is_match = all((ks == 1 for ks in kernel_size))
                else:
                    is_match = kernel_size == 1
            else:
                is_match = False

            if is_match:
                custom_data.append((node, kernel_size, stride))
                return True
            else:
                return False

        pool_one_kernel_size_nodes = []
        graph.filter_forward_nodes(_is_pool_nd_with_one_kernel_size, pool_one_kernel_size_nodes)
        for idx, (node, kernel_size, stride) in enumerate(pool_one_kernel_size_nodes):
            slices = [slice(None)] * 2
            t = node.prev_tensors[0]
            dim = len(t.shape)
            if not isinstance(stride, (list, tuple)):
                stride = [stride] * (dim - 2)
            for s in stride:
                if s == 1 or s is None:
                    slices.append(slice(None))
                else:
                    slices.append(slice(None, None, s))

            with override_current_trace_graph(graph):
                new_func = TraceFunction('torch.Tensor.__getitem__', True).parse_args(t, slices)

                graph.module_unique_name_dict[id(new_func)] = f'rewritten_pool_{idx}'
                graph.module_original_name_dict[id(new_func)] = f'rewritten_pool_{idx}'

                graph.replace_node_module(node, new_func)

        # Rewrite Linear-BatchNorm1d structure to Conv2d-BatchNorm2d
        is_rewrite_to_fuse = functools.partial(
            self.is_fusable,
            current_rules=load_processed_rewrite_to_fuse_rules(),
            check_node_quantized=False,
            graph=graph,
            layerwise_config_default=True,
            use_original_name=False,
        )
        custom_data = ([], set())
        graph.filter_forward_nodes(is_rewrite_to_fuse, custom_data, reverse=True)
        rewrite_fuse_names_list = custom_data[0]
        log.debug(f'found names_list that need to rewrite for fusing: {rewrite_fuse_names_list}')

        for idx, names in enumerate(reversed(rewrite_fuse_names_list)):
            # case fc-bn1d
            assert len(names) == 2, 'the rewrite nodes list length != 2'

            node_fc = graph.nodes_map[names[0]]
            node_bn1d = graph.nodes_map[names[1]]
            mod_fc = node_fc.module
            mod_bn = node_bn1d.module

            assert type(mod_fc) is nn.Linear and type(mod_bn) is nn.BatchNorm1d, "the rewrite struct is\'t [fc-bn1d]"

            if len(node_fc.prev_tensors[0].shape) != 2:
                log.debug('the [fc-bn]\'s input dimension != 2')
                continue
            # for fc-bn1d, rewrite [fc-bn1d] to [conv2d-bn2d]
            new_conv2d = torch.nn.Conv2d(
                in_channels=mod_fc.in_features,
                out_channels=mod_fc.out_features,
                kernel_size=[1, 1],
                bias=mod_fc.bias is not None,
            )
            fc_weight = mod_fc.weight
            new_conv2d.weight = nn.Parameter(torch.reshape(fc_weight, [fc_weight.shape[0], fc_weight.shape[1], 1, 1]))
            if mod_fc.bias is not None:
                new_conv2d.bias = mod_fc.bias
            graph.module_unique_name_dict[id(new_conv2d)] = f'rewritten_conv2d_bn2d_conv2d_{idx}'
            graph.module_original_name_dict[id(new_conv2d)] = f'rewritten_conv2d_bn2d_conv2d_{idx}'

            new_bn2d = torch.nn.BatchNorm2d(
                mod_bn.num_features,
                mod_bn.eps,
                mod_bn.momentum,
                affine=mod_bn.affine,
                track_running_stats=mod_bn.track_running_stats,
            )
            new_bn2d.load_state_dict(mod_bn.state_dict())
            graph.module_unique_name_dict[id(new_bn2d)] = f'rewritten_conv2d_bn2d_bn2d_{idx}'
            graph.module_original_name_dict[id(new_bn2d)] = f'rewritten_conv2d_bn2d_bn2d_{idx}'

            # replace new node, then insert reshape before new_conv2d and after new_bn2d
            with override_current_trace_graph(graph):
                graph.replace_node_module(node_fc, new_conv2d)
                graph.replace_node_module(node_bn1d, new_bn2d)

                prev_func = TraceFunction('torch.Tensor.__getitem__', prefix='rewritten_conv2d_bn2d_').parse_args(
                    node_fc.prev_tensors[0], (Ellipsis, None, None)
                )
                next_func = TraceFunction('torch.flatten', prefix='rewritten_conv2d_bn2d_').parse_args(
                    node_bn1d.next_tensors[0], 1
                )
            # expand the tensor shape between fc new_conv2d and new_bn2d
            node_fc.next_tensors[0].unsqueeze_(2).unsqueeze_(2)
            node_bn1d.prev_tensors[0].unsqueeze_(2).unsqueeze_(2)
            node_bn1d.next_tensors[0].unsqueeze_(2).unsqueeze_(2)

            prev_out = node_fc.prev_tensors[0][..., None, None]
            graph.insert_between(node_fc.prev_nodes[0], node_fc, prev_func, [prev_out], True)
            next_out = torch.flatten(node_bn1d.next_tensors[0], 1)
            graph.insert_after(node_bn1d, next_func, [next_out])

        # Rewrite BatchNorm1d to BatchNorm2d
        def _is_batch_norm_1d(node, custom_data):
            if self.layerwise_config.get(node.unique_name, True) is False:
                return False
            cur_module = node.module
            cur_class = type(cur_module)
            if len(node.prev_nodes) != 1:
                return False

            if node.prev_nodes[0].kind() in ('conv1d', nn.Conv1d):
                return False

            if cur_class == TraceFunction:
                return cur_module.kind == 'batch_norm' and node.prev_tensors[0].ndim == 3
            else:
                return cur_class == nn.BatchNorm1d

        batch_norm_1d_nodes = graph.filter_forward_nodes(_is_batch_norm_1d)
        for idx, node in enumerate(batch_norm_1d_nodes):
            mod = node.module
            if type(mod) is nn.BatchNorm1d:
                new_bn = torch.nn.BatchNorm2d(
                    mod.num_features,
                    mod.eps,
                    mod.momentum,
                    affine=mod.affine,
                    track_running_stats=mod.track_running_stats,
                )
                new_bn.load_state_dict(mod.state_dict())

                graph.module_unique_name_dict[id(new_bn)] = f'rewritten_bn2d_{idx}'
                graph.module_original_name_dict[id(new_bn)] = f'rewritten_bn2d_{idx}'

                with override_current_trace_graph(graph):
                    graph.replace_node_module(node, new_bn)

                    prev_func = TraceFunction('torch.unsqueeze', prefix='rewritten_bn2d_').parse_args(
                        node.prev_tensors[0], 2
                    )
                    next_func = TraceFunction('torch.squeeze', prefix='rewritten_bn2d_').parse_args(
                        node.next_tensors[0], 2
                    )

                node.next_tensors[0].unsqueeze_(2)

                prev_out = torch.unsqueeze(node.prev_tensors[0], 2)
                graph.insert_between(node.prev_nodes[0], node, prev_func, [prev_out], True)
                next_out = torch.squeeze(node.next_tensors[0], 2)
                graph.insert_after(node, next_func, [next_out])

        type_dict = {}
        for n in graph.forward_nodes:
            if n.type() not in (torch_q.QuantStub, torch_q.DeQuantStub):
                self.layerwise_config.setdefault(n.unique_name, True)
                kind = n.kind()
                type_str = kind if isinstance(kind, str) else kind.__name__
                type_dict[n.unique_name] = type_str

        fuse_mapping = {}
        if self.fused_layerwise_config:
            # Find all fusable nodes
            if type(self).__name__ == 'QATQuantizer':
                processed_all_rules = load_processed_all_qat_rules()
            else:
                processed_all_rules = load_processed_all_ptq_rules()

            is_fusable = functools.partial(
                self.is_fusable,
                current_rules=processed_all_rules,
                check_node_quantized=False,
                use_original_name=False,
            )

            custom_data = ([], set())
            graph.filter_forward_nodes(is_fusable, custom_data, reverse=True)
            activ_names = custom_data[0]

            for idx, names in enumerate(reversed(activ_names)):
                types = [graph.nodes_map[n].kind() for n in names]
                types = [x.__name__ if not isinstance(x, str) else x for x in types]
                types_str = ', '.join(types)
                types_str = f'({types_str})'
                type_dict[names[0]] = types_str
                for name in names[1:]:
                    fuse_mapping[name] = names[0]
                    self.layerwise_config[name] = self.layerwise_config.get(names[0], True)

        for n, t in type_dict.items():
            self.layerwise_config.yaml_add_eol_comment(f'type: {t}', n)

        skip_types = set(k[0] for k in REWRITE_QUANTIZABLE_RULE_LIST if len(k) == 1)
        for module_cls, action in self.quantize_op_action.items():
            if action == 'rewrite':
                skip_types.add(module_cls)
        if self.set_quantizable_op_stats:
            skip_types |= set(KNOWN_QSTATS.keys())
        skip_types_prev = skip_types | set(k[-1] for k in REWRITE_QUANTIZABLE_RULE_LIST if len(k) > 1)
        skip_types_next = skip_types | set(k[0] for k in REWRITE_QUANTIZABLE_RULE_LIST if len(k) > 1)

        # Add quant/dequant nodes for non-quantizable OPs
        disable_quantize_op_list = UNSUPPORTED_PYTORCH_QUANTIZATION_OP_LIST.copy()
        for module_cls, action in self.quantize_op_action.items():
            if action in ('disable', 'rewrite'):
                disable_quantize_op_list[module_cls] = None

        def _is_rewritable_lstm_node(node, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            return cur_class == nn.LSTM

        if self.quantize_op_action.get(nn.LSTM, 'enable') == 'rewrite':
            rewritable_lstm_nodes = graph.filter_forward_nodes(_is_rewritable_lstm_node)
            fake_dequant_cls = torch_q.DeQuantStub
            for idx, node in enumerate(rewritable_lstm_nodes):
                assert node.module.num_layers == 1, "Quantization rewrite for multi-layer LSTM is not yet supported"
                assert not node.module.bidirectional, "Quantization rewrite for bidirectional LSTM is not yet supported"

                cell_state = node.next_tensors[1][1]

                fake_dequant = fake_dequant_cls()

                fake_dequant_name = f'fake_dequant_rewrite_{idx}'

                graph.module_unique_name_dict[id(fake_dequant)] = fake_dequant_name
                graph.module_original_name_dict[id(fake_dequant)] = fake_dequant_name

                module_constructor_lines[id(fake_dequant)] = f'{qualified_name(fake_dequant_cls)}()'

                new_node = graph.insert_new_after(
                    node, fake_dequant, [cell_state], [[1, 1]], before_node=node.next_nodes[0]
                )

                with override_current_trace_graph(graph):
                    size_func = TraceFunction(
                        'torch.Tensor.size', is_class=True, prefix='fake_dequant_rewrite_'
                    ).parse_args(new_node.next_tensors[0], -1)

                size_node = graph.insert_new_after(
                    new_node,
                    size_func,
                    [new_node.next_tensors[0]],
                    [None],
                    next_tensors=[torch.tensor(new_node.next_tensors[0].size(-1))],
                    before_node=node.next_nodes[0],
                )
                size_len = len(node.next_tensors[0].shape)

                if node.module.bidirectional:
                    with override_current_trace_graph(graph):
                        size_func = TraceFunction(
                            'torch.Tensor.__mul__', is_class=True, prefix='fake_dequant_rewrite_'
                        ).parse_args(size_node.next_tensors[0], 2)

                        size_node = graph.insert_new_after(
                            size_node,
                            size_func,
                            [size_node.next_tensors[0]],
                            [None],
                            next_tensors=[size_node.next_tensors[0] * 2],
                            before_node=node.next_nodes[0],
                        )

                with override_current_trace_graph(graph):
                    expand_func = TraceFunction(
                        'torch.Tensor.expand', is_class=True, prefix='fake_dequant_rewrite_'
                    ).parse_args(node.next_tensors[0], *((-1,) * (size_len - 1)), size_node.next_tensors[0])

                graph.insert_between(
                    node, node.next_nodes[0], expand_func, tensor_ptrs=[id(node.next_tensors[0])], move_idx=True
                )
                expand_node = graph.nodes_map[expand_func.unique_name]
                size_node.next_nodes.append(expand_node)
                expand_node.prev_nodes.append(node)
                expand_node.prev_tensors.append(size_node.next_tensors[0])
                expand_node.prev_indices.append(None)

        def _is_not_quantizable(node, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == ConstantNode:
                return False
            elif cur_class == TraceFunction:
                if node.type() in ('__truediv__', '__itruediv__', 'div', 'div_'):
                    if node.prev_nodes[0].kind() in ('shape', 'size'):
                        return False
                if node.type() in ('shape', 'device', 'size', 'dtype'):
                    if node.unique_name in self.layerwise_config:
                        del self.layerwise_config[node.unique_name]
                    return False
                if self.layerwise_config.get(node.unique_name, True) is False:
                    return True
                supported_version = disable_quantize_op_list.get(cur_module.kind, torch.__version__)
                return supported_version is None or LooseVersion(torch.__version__) < supported_version
            else:
                if isinstance(cur_module, (torch_q.QuantStub, torch_q.DeQuantStub)):
                    return False
                if self.layerwise_config.get(node.unique_name, True) is False:
                    return True
                unsupported_types = tuple(
                    k
                    for k, v in disable_quantize_op_list.items()
                    if type(k) is not str
                    and k not in Q_MODULES_MAPPING
                    and (v is None or LooseVersion(torch.__version__) < v)
                )
                return isinstance(cur_module, unsupported_types)

        unsupported_nodes = graph.filter_forward_nodes(_is_not_quantizable)
        for idx, node in enumerate(reversed(unsupported_nodes)):
            if node.unique_name in self.layerwise_config:
                if node.kind() not in skip_types and self.layerwise_config[node.unique_name]:
                    del self.layerwise_config[node.unique_name]

            node_map = dict()
            next_nodes = {n.unique_name: n for n in node.next_nodes}.values()
            for inner_idx, next_node in enumerate(next_nodes):
                prev_tensor_ptrs = []
                if type(next_node.module) is TraceFunction and next_node.module.is_property:
                    continue

                for pt in next_node.prev_tensors:
                    for nt in node.next_tensors:
                        if isinstance(nt, (list, tuple)):
                            for k, ntt in enumerate(nt):
                                if id(pt) == id(ntt):
                                    if id(pt) not in prev_tensor_ptrs:
                                        prev_tensor_ptrs.append(id(pt))
                                    break
                        elif id(pt) == id(nt):
                            if id(pt) not in prev_tensor_ptrs:
                                prev_tensor_ptrs.append(id(pt))
                            break

                for ptr_idx, ptr in enumerate(prev_tensor_ptrs):
                    if ptr in node_map:
                        fake_quant = node_map[ptr]

                        graph.insert_between(node, next_node, fake_quant, move_idx=True, tensor_ptrs=[ptr])
                    else:
                        fake_quant = torch_q.QuantStub()

                        fake_quant_name = f'fake_quant_inner_{idx}_{inner_idx}_{ptr_idx}'

                        graph.module_unique_name_dict[id(fake_quant)] = fake_quant_name
                        graph.module_original_name_dict[id(fake_quant)] = fake_quant_name

                        fake_quant_cls = type(fake_quant)
                        module_constructor_lines[id(fake_quant)] = f'{qualified_name(fake_quant_cls)}()'

                        graph.insert_between(node, next_node, fake_quant, move_idx=True, tensor_ptrs=[ptr])
                        node_map[ptr] = graph.nodes_map[fake_quant_name]

        # Insert the DeQuantStub nodes before every input node of the unsupported ops
        for idx, node in enumerate(unsupported_nodes):
            fake_dequant_cls = torch_q.DeQuantStub
            assert node.rev_index is False
            node_map = dict()
            prev_nodes = {n.unique_name: n for n in node.prev_nodes}.values()
            for inner_idx, prev_node in enumerate(prev_nodes):

                if prev_node.kind() in ('shape', 'device', 'size', 'dtype'):
                    continue

                prev_tensor_ptrs = []
                for pt in node.prev_tensors:
                    for nt in prev_node.next_tensors:
                        if isinstance(nt, (list, tuple)):
                            for k, ntt in enumerate(nt):
                                if id(pt) == id(ntt):
                                    if id(pt) not in prev_tensor_ptrs:
                                        prev_tensor_ptrs.append(id(pt))
                                    break
                        elif id(pt) == id(nt):
                            if id(pt) not in prev_tensor_ptrs:
                                prev_tensor_ptrs.append(id(pt))
                            break

                for ptr_idx, ptr in enumerate(prev_tensor_ptrs):
                    if ptr in node_map:
                        fake_dequant = node_map[ptr]

                        graph.insert_between(prev_node, node, fake_dequant, move_idx=True, tensor_ptrs=set([ptr]))
                    else:

                        fake_dequant = fake_dequant_cls()

                        fake_dequant_name = f'fake_dequant_inner_{idx}_{inner_idx}_{ptr_idx}'

                        graph.module_unique_name_dict[id(fake_dequant)] = fake_dequant_name
                        graph.module_original_name_dict[id(fake_dequant)] = fake_dequant_name

                        module_constructor_lines[id(fake_dequant)] = f'{qualified_name(fake_dequant_cls)}()'

                        graph.insert_between(prev_node, node, fake_dequant, move_idx=True, tensor_ptrs=set([ptr]))
                        node_map[ptr] = graph.nodes_map[fake_dequant_name]

        # Remove consecutive dequant quant nodes
        def _is_consecutive_dequant_quant_nodes(node, custom_data):
            cur_type = node.type()
            if cur_type in (torch_q.QuantStub, torch_q.DeQuantStub):
                for next_node in node.next_nodes:
                    next_type = next_node.type()
                    if next_type in (torch_q.QuantStub, torch_q.DeQuantStub):
                        if cur_type != next_type:
                            if cur_type == torch_q.QuantStub:
                                target_node = None
                                if len(node.prev_nodes) == 1 and node.prev_nodes[0].kind() in skip_types_prev:
                                    target_node = node.prev_nodes[0]
                                elif (
                                    len(next_node.next_nodes) == 1 and next_node.next_nodes[0].kind() in skip_types_next
                                ):
                                    target_node = next_node.next_nodes[0]
                                if target_node is not None:
                                    self.layerwise_config.setdefault(target_node.unique_name, True)
                                    if self.layerwise_config[target_node.unique_name]:
                                        return False
                            custom_data.append((node, next_node))
                            return True
            return False

        consecutive_dequant_quant_nodes = []
        consecutive_branch_dequant_quant_nodes = []
        graph.filter_forward_nodes(_is_consecutive_dequant_quant_nodes, consecutive_dequant_quant_nodes)
        for node, next_node in consecutive_dequant_quant_nodes:
            if len(node.next_nodes) == 1 and len(next_node.prev_nodes) == 1:
                graph.remove_node(next_node)
                graph.remove_node(node)
            else:
                consecutive_branch_dequant_quant_nodes.append((node, next_node))

        # TODO: Support complex cases
        for node, next_node in consecutive_branch_dequant_quant_nodes:
            is_removable = all(
                (
                    n.type() in (torch_q.QuantStub, torch_q.DeQuantStub, 'shape', 'device', 'size', 'dtype')
                    and n.type() != node.type()
                    for n in node.next_nodes
                )
            )
            if is_removable:
                graph.remove_node(node)
                for n in node.next_nodes:
                    if n.type() in (torch_q.QuantStub, torch_q.DeQuantStub):
                        graph.remove_node(n)

        # Process additional fusable nodes
        processed_extra_q_rules = load_processed_extra_q_rules()
        is_extra_fusable = functools.partial(
            self.is_fusable,
            current_rules=processed_extra_q_rules,
            check_node_quantized=False,
            use_original_name=False,
        )

        custom_data = ([], set())
        graph.filter_forward_nodes(is_extra_fusable, custom_data, reverse=True)
        activ_names = custom_data[0]
        log.debug(f'found nodes that cannot fuse: {activ_names}')

        for idx, names in enumerate(reversed(activ_names)):
            name = names[-1]
            node = graph.nodes_map[name]
            fake_quant = torch_q.QuantStub()

            graph.module_unique_name_dict[id(fake_quant)] = f'fake_activ_quant_{idx}'
            graph.module_original_name_dict[id(fake_quant)] = f'fake_activ_quant_{idx}'

            fake_quant_cls = type(fake_quant)
            module_constructor_lines[id(fake_quant)] = f'{qualified_name(fake_quant_cls)}()'

            graph.insert_after(node, fake_quant)

            fake_dequant = torch_q.DeQuantStub()

            graph.module_unique_name_dict[id(fake_dequant)] = f'fake_activ_dequant_{idx}'
            graph.module_original_name_dict[id(fake_dequant)] = f'fake_activ_dequant_{idx}'

            fake_dequant_cls = type(fake_dequant)
            module_constructor_lines[id(fake_dequant)] = f'{qualified_name(fake_dequant_cls)}()'

            graph.insert_after(node, fake_dequant)

        fused_clamps = []
        for names in activ_names:
            for name in names:
                node = graph.nodes_map[name]
                if node.kind() == 'clamp_with_fusion':
                    fused_clamps.append(name)

        # Rewrite unfused clamp_with_fusion to torch.clamp
        def _is_unfused_clamp_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return (
                    cur_module.kind == 'clamp_with_fusion'
                    and torch.is_floating_point(node.next_tensors[0])
                    and node.unique_name not in fused_clamps
                )
            return False

        unfused_clamp_nodes = graph.filter_forward_nodes(_is_unfused_clamp_node)
        log.info(f'rewriting unfused_clamp for {[node.unique_name for node in unfused_clamp_nodes]}')
        for node in unfused_clamp_nodes:
            node.module.kind = kind
            node.module.func_type = node.module.func_type.replace('clamp_with_fusion', 'clamp')
            node.module.full_name = f'torch.{node.module.func_type}'

        # Optional tensor shape broadcasting for quantized binary ops
        def _is_broadcastable_binary_quantized_op_node(node: TraceNode, custom_data) -> bool:
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class != TraceFunction:
                return False
            return (
                cur_module.full_name.startswith('self.float_functional_simple_')
                and cur_module.func_type in ('add', 'mul', 'add_relu')
                and node.prev_tensors[0].shape != node.prev_tensors[1].shape
            )

        broadcastable_binary_quantized_op_nodes = graph.filter_forward_nodes(_is_broadcastable_binary_quantized_op_node)
        for node in broadcastable_binary_quantized_op_nodes:
            assert len(node.prev_nodes) == 2
            assert len(node.prev_tensors) == 2

            l_shape = list(node.prev_tensors[0].shape)
            r_shape = list(node.prev_tensors[1].shape)

            ref_index = None
            if len(l_shape) > len(r_shape):
                ref_index = 0
                r_shape = [1] * (len(l_shape) - len(r_shape)) + r_shape
            elif len(l_shape) < len(r_shape):
                ref_index = 1
                l_shape = [1] * (len(r_shape) - len(l_shape)) + l_shape

            for l_dim, r_dim in zip(l_shape, r_shape):
                if l_dim > r_dim:
                    if ref_index in (None, 0) and r_dim == 1:
                        ref_index = 0
                    else:
                        ref_index = -1
                        break
                elif l_dim < r_dim:
                    if ref_index in (None, 1) and l_dim == 1:
                        ref_index = 1
                    else:
                        ref_index = -1
                        break

            if ref_index >= 0:
                src_index = 1 - ref_index

                with override_current_trace_graph(graph):
                    trace_func = TraceFunction('torch.Tensor.expand_as', True, prefix='fuse_').parse_args(
                        node.prev_tensors[src_index], node.prev_tensors[ref_index]
                    )
                next_tensors = [node.prev_tensors[src_index].expand_as(node.prev_tensors[ref_index])]
                graph.insert_between(node.prev_nodes[src_index], node, trace_func, next_tensors, True)

                new_node = graph.nodes_map[trace_func.unique_name]
                new_node.prev_nodes.append(node.prev_nodes[ref_index])
                new_node.prev_tensors.append(node.prev_tensors[ref_index])
                new_node.prev_indices.append(node.prev_indices[ref_index])
                node.prev_nodes[ref_index].next_nodes.append(new_node)
            else:
                new_indices = []
                for i in range(2):
                    if node.prev_indices[i] is None:
                        new_indices.append(i)
                    elif isinstance(node.prev_indices[i], (tuple, list)):
                        new_indices.append(node.prev_indices[i] + [i])
                    else:
                        new_indices.append([node.prev_indices[i], i])
                with override_current_trace_graph(graph):
                    trace_func = TraceFunction('torch.broadcast_tensors', False, prefix='fuse_').parse_args(
                        node.prev_tensors[0], node.prev_tensors[1]
                    )
                next_tensors = torch.broadcast_tensors(node.prev_tensors[0], node.prev_tensors[1])
                graph.insert_before(node, trace_func, next_tensors, False, new_indices)

        for name in fuse_mapping:
            if name in self.layerwise_config:
                del self.layerwise_config[name]

        for name in self.layerwise_config:
            node = graph.nodes_map[name]
            if isinstance(node.module, nn.Module) and not isinstance(
                node.module,
                (nn.Dropout, nn.AdaptiveAvgPool2d, nn.AvgPool2d, nn.MaxPool2d, nn.ReLU, nn.Upsample, nn.ConstantPad2d),
            ):
                self.effective_layers.append(name)
            elif node.kind() in ('add', 'mul', 'cat') or node.kind() in FUNCTIONAL_MODULE_MAPPING:
                self.effective_layers.append(name)

        graph.quantized = True
        graph.recompute_forward_order()

    def is_fusable(
        self,
        node,
        custom_data,
        current_rules=None,
        check_node_quantized=True,
        use_original_name=True,
        graph=None,
        layerwise_config_default=False,
    ):
        # Tell whether a TraceNode is fusable with some nearby nodes

        if current_rules is None:
            return False

        cur_node = node
        names = []
        final_names = []
        current_state = False
        while True:
            if check_node_quantized:
                if not cur_node.quantized:
                    break
            else:
                if self.layerwise_config.get(cur_node.unique_name, layerwise_config_default) is False:
                    break
            cur_module = cur_node.module
            cur_class = type(cur_module)
            if isinstance(cur_module, TraceFunction):
                cur_class = cur_module.kind
            prev_nodes = cur_node.prev_nodes
            log.debug(f'cur: {cur_class}')
            if cur_class in current_rules or cur_class == nn.Identity:
                if use_original_name and isinstance(cur_module, nn.Module):
                    cur_name = graph.module_original_name_dict[id(cur_module)]
                else:
                    cur_name = cur_node.unique_name
                if cur_name in custom_data[1]:
                    log.debug('found existing nodes, skipping')
                    break
                if len(prev_nodes) == 0:
                    break
                if cur_class in current_rules:
                    current_state, current_rules = current_rules[cur_class]
                    log.debug(f'dict: {current_rules}, {current_state}')
                    names.append(cur_name)
                    if current_state is True:
                        log.debug(f'update best: {names}')
                        final_names.clear()
                        final_names.extend(names)
                if len(cur_node.prev_nodes) != 1:
                    break
                cur_node = cur_node.prev_nodes[0]
            else:
                if cur_class == nn.ConvTranspose2d:
                    if len(final_names) > 0:
                        next_types = []
                        for _ in range(len(final_names)):
                            next_node = cur_node.next_nodes[0]
                            next_types.append(next_node.kind())
                            cur_node = next_node
                        if tuple(next_types) == (nn.BatchNorm2d, nn.ReLU):
                            final_names.clear()
                break

        if len(final_names) > 0:
            final_names.reverse()
            log.debug(f'final: {final_names}')
            custom_data[0].append(final_names)
            for name in final_names:
                custom_data[1].add(name)
            return True

        return False

    def error_analysis(self, qat_model: nn.Module, dummy_input, threshold: float = 20.0):
        """Generates the QAT error report using the SQNR metric

        Args:
            qat_model: The QAT model
            dummy_input: A viable input to the model
            threshold (float): The threshold of SQNR. Defaults to 20.0
        """

        if isinstance(qat_model, DataParallel) or isinstance(qat_model, DistributedDataParallel):
            model = qat_model.module
        else:
            model = qat_model

        modules_list = {}
        names_list = {}

        float_results = {}
        hooks = []

        def forward_hook(module, input, output):
            name = names_list[module]
            float_results[name] = input

        fake_quant_enabled_dict = {}
        observer_enabled_dict = {}
        for n, m in model.named_modules():
            if isinstance(m, torch.quantization.FakeQuantize):
                names_list[m] = n
                modules_list[n] = m

                fake_quant_enabled_dict[m] = m.fake_quant_enabled.clone()
                observer_enabled_dict[m] = m.observer_enabled.clone()

                hooks.append(m.register_forward_hook(forward_hook))

        if len(modules_list) == 0:
            log.warning('No FakeQuantize modules found. Are you sure you passed in a QAT model?')
            return

        model.apply(torch.quantization.disable_fake_quant)
        model.apply(torch.quantization.disable_observer)

        device = get_module_device(model)

        if type(dummy_input) is torch.Tensor:
            actual_input = [dummy_input]
        elif isinstance(dummy_input, (tuple, list)):
            actual_input = list(dummy_input)
        else:
            log.error(f'Unsupported type {type(dummy_input)} for dummy input')
            assert False

        for i in range(len(actual_input)):
            dummy_input = actual_input[i]
            if type(dummy_input) is torch.Tensor:
                if dummy_input.device != device:
                    actual_input[i] = dummy_input.to(device)

        model.eval()

        with torch.no_grad():
            model(*actual_input)

        for h in hooks:
            h.remove()
        hooks.clear()

        for m, v in fake_quant_enabled_dict.items():
            m.fake_quant_enabled = v

        def sqnr(x, y):
            Ps = torch.norm(x)
            Pn = torch.norm(x - y)
            return (20 * torch.log10(Ps / Pn)).item()

        q_errors_weight = []
        q_errors_activ = []
        while len(float_results) > 0:
            n, f = float_results.popitem()
            mod = modules_list[n]
            with torch.no_grad():
                q = mod(*f)
                loss = sqnr(f[0], q)
            actual_n = '.'.join(n.split('.')[:-1])
            if loss <= threshold:
                if n.endswith('.weight_fake_quant'):
                    q_errors_weight.append((actual_n, mod, loss))
                else:
                    q_errors_activ.append((actual_n, mod, loss))

        q_errors_weight = sorted(q_errors_weight, key=lambda x: x[2])
        q_errors_activ = sorted(q_errors_activ, key=lambda x: x[2])

        logs = []
        if len(q_errors_weight) > 0:
            logs.append('')
            logs.append(f'Weights (SQNR <= {threshold}):')
            for n, m, e in q_errors_weight:
                logs.append(f'{n} SQNR: {e:.4f}, scale: {m.scale.item():.4f}, zero_point: {m.zero_point.item()}')

        if len(q_errors_activ) > 0:
            logs.append('')
            logs.append(f'Activations (SQNR <= {threshold}):')
            for n, m, e in q_errors_activ:
                logs.append(f'{n} SQNR: {e:.4f}, scale: {m.scale.item():.4f}, zero_point: {m.zero_point.item()}')

        if len(q_errors_weight) == 0 and len(q_errors_activ) == 0:
            logs.append('')
            logs.append('All good!')

        if len(logs) > 0:
            logs.insert(0, 'Quantization error report:')
            logs.append('')

            full_log = '\n'.join(logs)
            log.warning(full_log)

        for m, v in observer_enabled_dict.items():
            m.observer_enabled = v

        model.train()

    def restore_to_original(self, q_model: nn.Module):
        """Restores a QAT/PTQ-converted model to original state

        Args:
            qat_model: The QAT/PTQ-converted model

        """

        sub_list = []
        for n, m in q_model.named_children():
            if isinstance(m, (torch.nn.quantized.Quantize, torch.nn.quantized.DeQuantize)):
                sub_list.append((n, nn.Identity()))
            elif isinstance(m, torch.nn.quantized.Linear):
                fc = nn.Linear(m.in_features, m.out_features, m.bias is not None)
                fc.weight = torch.nn.Parameter(m.weight().dequantize())
                fc.bias = torch.nn.Parameter(m.bias())
                sub_list.append((n, fc))
            elif isinstance(m, torch.nn.quantized.Conv2d):
                conv = nn.Conv2d(
                    m.in_channels,
                    m.out_channels,
                    m.kernel_size,
                    m.stride,
                    m.padding,
                    m.dilation,
                    m.groups,
                    m.bias is not None,
                    m.padding_mode,
                )
                conv.weight = torch.nn.Parameter(m.weight().dequantize())
                conv.bias = torch.nn.Parameter(m.bias())
                sub_list.append((n, conv))
            elif isinstance(m, torch.nn.quantized.QFunctional):
                sub_list.append((n, nn.quantized.FloatFunctional()))
            elif hasattr(torch.nn.quantized, 'ReLU') and isinstance(m, torch.nn.quantized.ReLU):
                sub_list.append((n, nn.ReLU()))
            elif isinstance(m, torch.nn.quantized.ReLU6):
                sub_list.append((n, nn.ReLU6()))
            elif type(m).__module__.startswith('torch.nn.modules'):
                continue
            else:
                assert False, f"unsupported type: {type(m).__name__}"

        for n, m in sub_list:
            setattr(q_model, n, m)

    def convert(self, q_model: nn.Module, backend: str = 'tflite') -> nn.Module:
        """Converts a QAT/PTQ-prepared model to an actual quantized model

        Args:
            q_model (nn.Module): The QAT/PTQ-prepared model
            backend (str): The backend to translate for, including `pytorch` and `tflite`. Defaults to `tflite`

        Returns:
            nn.Module: The QAT/PTQ-converted model. When the backend is set to `pytorch`, it is used for validation \
                in PyTorch only.
        """
        if self.quantize_op_action.get(nn.LSTM, None) and self.backend == 'qnnpack':
            self.restore_lstm_weights(q_model)

        for acp, post_acp, dq_name, q_name, activ_name, activ_type in self.extra_qparams_mappings:
            if backend != 'pytorch' and activ_type in ('relu', 'relu6', torch.nn.ReLU, torch.nn.ReLU6):
                continue

            acp.scale = post_acp.scale
            acp.zero_point = post_acp.zero_point
            acp.activation_post_process.min_val = post_acp.activation_post_process.min_val
            acp.activation_post_process.max_val = post_acp.activation_post_process.max_val

            setattr(q_model, dq_name, nn.Identity())
            setattr(q_model, q_name, nn.Identity())
            if activ_name is not None:
                setattr(q_model, activ_name, nn.Identity())

        if type(self).__name__ == 'QATQuantizer':

            q = queue.Queue()
            q.put(q_model)

            while not q.empty():
                m = q.get()
                for n, c in m.named_children():
                    if isinstance(c, ConvTransposeBn2d):
                        setattr(m, n, c.transform(c))
                    else:
                        q.put(c)

            if hasattr(torch_q, 'get_default_static_quant_module_mappings'):
                mapping = torch_q.get_default_static_quant_module_mappings()
            elif hasattr(torch_q, 'get_static_quant_module_mappings'):
                mapping = copy.deepcopy(torch_q.get_static_quant_module_mappings())
            else:
                mapping = copy.deepcopy(torch_q.DEFAULT_MODULE_MAPPING)

            mapping.update(FUSE_QAT_MODULES_CVT)

            float_mods = {}

            for qat_t, q_t in FUSE_QAT_MODULES_CVT.items():
                float_mod = getattr(q_t, '_FLOAT_MODULE', None)
                if float_mod is not None:
                    float_mods[q_t] = float_mod
                    setattr(q_t, '_FLOAT_MODULE', qat_t)

            q_model = torch.quantization.convert(q_model, mapping)

            for q_t, orig_t in float_mods.items():
                setattr(q_t, '_FLOAT_MODULE', orig_t)

            float_mods.clear()
        else:
            q_model = torch.quantization.convert(q_model)

        return q_model

    def optimize_conv_bn_fusion(self, q_model, eps=1e-5):
        """Optimizes the Conv-BatchNorm fusion pattern.
           Sometimes, the running_var of the BatchNorm could be near to zero. If the weights of those channels happen
           to be large, then it may lead to large quantization losses. We choose to ignore those channels when
           `bn.running_var.abs() < eps`.

        Args:
            q_model (nn.Module): The QAT/PTQ-prepared model

        """

        def _pre_hook_func(indices):
            def _pre_hook(mod, input):
                max_val = input[0][~indices].max()
                min_val = input[0][~indices].min()
                input[0].clamp_(min_val, max_val)
                return input

            return _pre_hook

        for m in q_model.modules():
            if type(m).__name__ in ('ConvBnReLU2d', 'ConvBn2d'):
                if m.in_channels == m.out_channels and m.out_channels == m.groups and m.groups > 1:
                    indices = m.bn.running_var < eps
                    if torch.any(indices):
                        m.weight_fake_quant.register_forward_pre_hook(_pre_hook_func(indices))

    def prepare_onnx_export(self, q_model):
        """Prepares for ONNX model export

        Args:
            q_model (nn.Module): The QAT/PTQ-prepared model

        """

        if self.backend == 'onnx':
            for mod in self.leaf_nodes:
                torch.quantization.disable_fake_quant(mod.activation_post_process)

        fq_base_cls = getattr(torch_q, 'FakeQuantizeBase', torch_q.FakeQuantize)
        for n, m in q_model.named_modules():
            if isinstance(m, fq_base_cls):
                if m.qscheme in (torch.per_channel_symmetric, torch.per_tensor_symmetric):
                    m.zero_point.fill_(0)
                else:
                    assert m.qscheme == torch.per_tensor_affine
                    assert isinstance(m, torch_q.FixedQParamsFakeQuantize)

                    m.qscheme = torch.per_tensor_symmetric
                    if torch.all(m.zero_point == 128):
                        m.zero_point.fill_(0)
                    else:
                        assert not torch.is_nonzero(m.zero_point)
                        m.scale *= 2

                m.quant_min = -128
                m.quant_max = 127
                m.dtype = torch.int8

                if hasattr(m, 'activation_post_process'):
                    m.activation_post_process.quant_min = -128
                    m.activation_post_process.quant_max = 127
                    m.activation_post_process.dtype = torch.int8

        if LooseVersion(torch.__version__) >= LooseVersion('1.12.0'):

            def get_wrapper(mod):
                def wrapper(*args, **kwargs):
                    return torch_q.FakeQuantize.forward(mod, *args, **kwargs)

                return wrapper

            fused_fq_cls = getattr(torch_q, 'FusedMovingAvgObsFakeQuantize', None)
            if fused_fq_cls is not None:
                for m in q_model.modules():
                    if isinstance(m, fused_fq_cls):
                        m.forward = get_wrapper(m)
        elif LooseVersion(torch.__version__) >= LooseVersion('1.10.0'):
            mod_dict = {}
            for n, m in q_model.named_modules():
                mod_dict[n] = m

            class _FakeQuantize(nn.Module):
                def __init__(self, fq: torch_q.FakeQuantize) -> None:
                    super().__init__()
                    self.fake_quant_enabled = fq.fake_quant_enabled
                    self.scale = fq.scale
                    self.zero_point = fq.zero_point
                    self.quant_min = fq.quant_min
                    self.quant_max = fq.quant_max
                    self.is_per_channel = fq.is_per_channel
                    self.ch_axis = fq.ch_axis

                def forward(self, X):
                    if self.fake_quant_enabled[0] == 1:
                        if self.is_per_channel:
                            X = torch.fake_quantize_per_channel_affine(
                                X, self.scale, self.zero_point, self.ch_axis, self.quant_min, self.quant_max
                            )
                        else:
                            X = torch.fake_quantize_per_tensor_affine(
                                X, self.scale, self.zero_point, self.quant_min, self.quant_max
                            )
                    return X

            action_list = []
            for n, m in q_model.named_modules():
                if isinstance(m, torch_q.FakeQuantize):
                    if '.' in n:
                        n_split = n.split('.')
                        prop = n_split[-1]
                        pm_name = '.'.join(n_split[:-1])
                        pm = mod_dict[pm_name]
                    else:
                        prop = n
                        pm = q_model

                    new_m = _FakeQuantize(m)

                    action_list.append((pm, prop, new_m))

            for parent, prop, new_mod in action_list:
                setattr(parent, prop, new_mod)

        for n, m in q_model.named_modules():

            def conv_fused_wrapper(mod, scale_factor):
                type_name = type(mod).__name__

                def new_wrapper(input):
                    weight_shape = [1] * len(mod.weight.shape)
                    weight_shape[0] = -1
                    y = type(mod)._conv_forward(
                        mod, input, mod.weight_fake_quant(mod.weight * scale_factor.reshape(weight_shape)), mod.bias
                    )
                    if 'Bn' in type_name:
                        y = mod.bn(y)
                    if 'ReLU' in type_name:
                        y = torch.nn.functional.relu(y)
                    return y

                return new_wrapper

            type_name = type(m).__name__
            if type_name in ('ConvBn1d', 'ConvBnReLU1d', 'ConvBn2d', 'ConvBnReLU2d', 'ConvBn3d', 'ConvBnReLU3d'):
                running_std = torch.sqrt(m.bn.running_var + m.bn.eps)
                scale_factor = m.bn.weight / running_std
                if m.bias is not None:
                    m.bn.running_mean -= m.bias
                    m.bias = None
                m.bn.running_mean *= scale_factor
                m.bn.weight /= scale_factor
                m.forward = conv_fused_wrapper(m, scale_factor)

        def get_pre_hook(acps, indices):
            def pre_hook(module, input):
                new_input = list(input)
                for acp, idx in zip(acps, indices):
                    new_input[idx] = acp(new_input[idx])
                return tuple(new_input)

            return pre_hook

        def get_hook_func(acps, indices, orig_func):
            def pre_hook(*input, **kwargs):
                new_input = list(input)
                for acp, idx in zip(acps, indices):
                    if orig_func.__name__ == 'cat':
                        new_input[0][idx] = acp(new_input[0][idx])
                    else:
                        new_input[idx] = acp(new_input[idx])
                input = tuple(new_input)
                return orig_func(*input, **kwargs)

            return pre_hook

        if self.backend == 'onnx':
            end_mod_dict = {}
            for start_mod, end_mod, idx in self.swap_nodes:
                end_mod_dict.setdefault(end_mod, ([], []))
                end_mod_dict[end_mod][0].append(start_mod)
                end_mod_dict[end_mod][1].append(idx)

            for end_mod, (start_mods, indices) in end_mod_dict.items():
                acps = []
                for start_mod in start_mods:
                    if inspect.isroutine(start_mod) and isinstance(start_mod.__self__, nnq.FloatFunctional):
                        acp = start_mod.__self__.activation_post_process
                    else:
                        acp = start_mod.activation_post_process
                    acps.append(acp)

                if inspect.isroutine(end_mod) and isinstance(end_mod.__self__, nnq.FloatFunctional):
                    ff = end_mod.__self__

                    for i in range(len(indices)):
                        indices[i] -= 1

                    ff.cat = get_hook_func(acps, indices, ff.cat)
                    ff.add = get_hook_func(acps, indices, ff.add)
                    ff.mul = get_hook_func(acps, indices, ff.mul)
                    ff.add_scalar = get_hook_func(acps, indices, ff.add_scalar)
                    ff.mul_scalar = get_hook_func(acps, indices, ff.mul_scalar)
                    ff.add_relu = get_hook_func(acps, indices, ff.add_relu)
                else:
                    assert isinstance(
                        end_mod, nn.Module
                    ), "Only end nodes with `nn.Module` are supported duing module swapping"

                    end_mod.register_forward_pre_hook(get_pre_hook(acps, indices))

        q_model.apply(torch_q.disable_observer)

    def freeze_weights(self, q_model, dummy_input):
        """Freezes the weights before model export without `quantizer.convert`

        Args:
            q_model (nn.Module): The QAT/PTQ-prepared model

        """
        hooks = []

        def collect_fake_quantize_hook(mod, inp, outp):
            mod.register_parameter('_tinynn_frozen_q_weight', nn.Parameter(outp))

        for n, m in q_model.named_modules():
            if n.endswith('.weight_fake_quant'):
                hooks.append(m.register_forward_hook(collect_fake_quantize_hook))

        if isinstance(dummy_input, (list, tuple)):
            q_model(*dummy_input)
        else:
            q_model(dummy_input)

        for hook in hooks:
            hook.remove()

        hooks.clear()

        def freeze_fake_quantize_hook(mod, inp):
            return mod._tinynn_frozen_q_weight

        for n, m in q_model.named_modules():
            if n.endswith('.weight_fake_quant'):
                hooks.append(m.register_forward_pre_hook(freeze_fake_quantize_hook))

    def fake_quantize_lstm_weights(self, asym=False, eps=1e-6):
        def _lstm_weight_fake_quantize(weight, asym=False, eps=1e-6):
            if weight.numel() == 0:
                return weight
            quant_min, quant_max = -127, 127
            weight_parts = torch.chunk(weight, 4)
            weight_quant_parts = []
            for i in range(4):
                weight_quant_parts.append(
                    fake_quantize(weight_parts[i], asym=asym, eps=eps, quant_min=quant_min, quant_max=quant_max)
                )
            weight_quant = torch.cat(weight_quant_parts)
            return weight_quant

        with torch.no_grad():
            for name, module in self.model.named_modules():
                if isinstance(module, torch.nn.LSTM):
                    for weight_name in ['weight_ih_l0', 'weight_hh_l0']:
                        quantized_weight = _lstm_weight_fake_quantize(getattr(module, weight_name), asym=asym, eps=eps)
                        self.lstm_origin_weight_dict[f"{name}.{weight_name}"] = getattr(module, weight_name).clone()
                        getattr(module, weight_name).data.copy_(quantized_weight)

    def restore_lstm_weights(self, model):
        with torch.no_grad():
            for name, module in model.named_modules():
                if isinstance(module, torch.nn.LSTM):
                    for weight_name in ['weight_ih_l0', 'weight_hh_l0']:
                        full_weight_name = f"{name}.{weight_name}"
                        if full_weight_name in self.lstm_origin_weight_dict:
                            getattr(module, weight_name).data.copy_(self.lstm_origin_weight_dict[full_weight_name])
        self.lstm_origin_weight_dict.clear()


class BF16Quantizer(QATQuantizer):
    def __init__(self, model, dummy_input, work_dir: typing.Optional[str] = None, config: typing.Optional[dict] = None):
        """ Constructs a new BF16Quantizer object

        Args:
            model: The model to be quantized
            dummy_input: A viable input to the model
            work_dir (typing.Optional[str], optional): The working directory in which the intermediate files will be \
                generated. Defaults to None, in which case "output" will be used.
            config (typing.Optional[dict]): Options for the quantizer
        """

        super().__init__(model, dummy_input, work_dir, config)

    def parse_config(self, config: dict):
        super().parse_config(config)

        self.rewrite_graph = False

    def quantize(self) -> nn.Module:
        """Prepare model for BFloat16 training"""

        self.model.train()

        qconfig = torch_q.QConfig(activation=FakeQuantizeBFloat16.with_args(), weight=FakeQuantizeBFloat16.with_args())
        self.model.qconfig = qconfig

        torch_q.prepare_qat(self.model, inplace=True)

        return self.model


class PostQuantizer(QATQuantizer):
    rewrite_graph: bool
    force_overwrite: bool
    is_input_quantized: typing.Optional[typing.Tuple[bool]]
    quantized_input_stats: typing.Optional[typing.List[typing.Optional[typing.Tuple[float, float]]]]
    backend: str
    remove_weights_after_load: bool
    asymmetric: bool
    per_tensor: bool
    disable_requantization_for_cat: bool
    dynamic_lstm_quant: bool
    algorithm: str

    def __init__(self, model, dummy_input, work_dir: typing.Optional[str] = None, config: typing.Optional[dict] = None):
        """ Constructs a new PostQuantizer object

        Args:
            model: The model to be quantized
            dummy_input: A viable input to the model
            work_dir (typing.Optional[str], optional): The working directory in which the intermediate files will be \
                generated. Defaults to None, in which case "output" will be used.
            config (typing.Optional[dict]): Options for the quantizer
        """

        super().__init__(model, dummy_input, work_dir, config)

    def parse_config(self, config: typing.Optional[dict]):
        if config and 'algorithm' not in config.keys():
            config['algorithm'] = 'l2'

        super().parse_config(config)

    def set_model_mode(self, model: nn.Module = None, preserve: bool = False):
        """Sets the mode (train/eval) of the model

        Args:
            model (nn.Module, optional): A PyTorch model. Defaults to None.
            preserve (bool, optional): Whether the old state is preserved. Defaults to False.
        """
        if model is None:
            model = self.model

        if preserve:
            for m in model.modules():
                training = getattr(m, 'training', None)
                if isinstance(training, bool):
                    self.train_mode_dict[m] = training

        model.eval()

    def prepare_qconfig(self, graph: TraceGraph, backend: str):
        """Prepare qconfig for various configurations.

        Args:
            graph (TraceGraph): The computation graph of the model
            backend (str, optional): The backend of quantization
        """

        log.info('setting qat backend and call prepare_qat')
        if not self.legacy_fq or LooseVersion(torch.__version__) < '1.12.0':
            qconfig = torch_q.get_default_qconfig(backend)
        else:
            qconfig = torch_q.get_default_qconfig(backend, 0)
        qconfig_c = None
        if self.backend == 'qnnpack':
            if not self.asymmetric:
                sym_fq = torch_q.HistogramObserver.with_args(
                    dtype=torch.quint8, qscheme=torch.per_tensor_symmetric, reduce_range=False
                )
                qconfig = torch_q.QConfig(sym_fq, qconfig.weight)
            if not self.per_tensor:
                sym_fq = MinMaxObserver.with_args(
                    dtype=torch.qint8, qscheme=torch.per_tensor_symmetric, reduce_range=False
                )
                qconfig = torch_q.QConfig(qconfig.activation, sym_fq)
                sym_fq = PerChannelMinMaxObserver.with_args(
                    dtype=torch.qint8, qscheme=torch.per_channel_symmetric, reduce_range=False, ch_axis=0
                )
                qconfig_c = torch_q.QConfig(qconfig.activation, sym_fq)
        elif self.backend == 'fbgemm':
            sym_fq = torch_q.MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric)
            qconfig_c = torch_q.QConfig(qconfig.activation, sym_fq)
        elif self.backend in ('onnx', 'tensorrt'):
            if not self.asymmetric:
                sym_fq = torch_q.HistogramObserver.with_args(
                    dtype=torch.qint8, qscheme=torch.per_tensor_symmetric, reduce_range=False
                )
                qconfig = torch_q.QConfig(sym_fq, qconfig.weight)
            if not self.per_tensor:
                sym_fq = torch_q.MinMaxObserver.with_args(
                    dtype=torch.qint8, qscheme=torch.per_tensor_symmetric, reduce_range=False
                )
                qconfig = torch_q.QConfig(qconfig.activation, sym_fq)
                sym_fq = torch_q.PerChannelMinMaxObserver.with_args(
                    dtype=torch.qint8, qscheme=torch.per_channel_symmetric, reduce_range=False, ch_axis=0
                )
                qconfig_c = torch_q.QConfig(qconfig.activation, sym_fq)
        else:
            log.warning(f'Quantization backend {self.backend} is not tested. Please use at your risk.')

        if self.algorithm != 'l2':
            if self.algorithm == 'kl':
                if self.backend == 'qnnpack':
                    alg_sym_fq = HistogramObserverKL.with_args(qscheme=torch.per_tensor_symmetric, reduce_range=False)
                    alg_asym_fq = HistogramObserverKL.with_args(reduce_range=False)
                elif self.backend == 'fbgemm':
                    alg_sym_fq = HistogramObserverKL.with_args(qscheme=torch.per_tensor_symmetric, reduce_range=True)
                    alg_asym_fq = HistogramObserverKL.with_args(reduce_range=True)
                else:
                    alg_sym_fq = qconfig.activation
                    alg_asym_fq = qconfig.activation
                if not self.asymmetric:
                    qconfig = torch_q.QConfig(alg_sym_fq, qconfig.weight)
                else:
                    qconfig = torch_q.QConfig(alg_asym_fq, qconfig.weight)

        torch.backends.quantized.engine = backend
        graph.module.qconfig = qconfig
        if self.backend == 'qnnpack':
            if qconfig_c is not None:
                q = queue.Queue()
                q.put(graph.module)

                while not q.empty():
                    m = q.get()
                    if type(m).__name__ in (
                        'Conv2d',
                        'ConvBnReLU2d',
                        'ConvBn2d',
                        'ConvReLU2d',
                        'Conv1d',
                        'ConvBnReLU1d',
                        'ConvBn1d',
                    ):
                        m.qconfig = qconfig_c
                    else:
                        for c in m.children():
                            q.put(c)
        elif self.backend == 'fbgemm':
            if qconfig_c is not None:
                q = queue.Queue()
                q.put(graph.module)

                while not q.empty():
                    m = q.get()
                    if type(m).__name__ in ('Linear', 'LinearReLU'):
                        m.qconfig = qconfig_c
                    else:
                        for c in m.children():
                            q.put(c)

        def _lstm_node(node, custom_data):
            return isinstance(node.module, nn.LSTM)

        if self.dynamic_lstm_quant:
            lstm_nodes = graph.filter_forward_nodes(_lstm_node)
            for node in lstm_nodes:
                node.quantized = True
                node.module.qconfig = torch_q.default_dynamic_qconfig

    def prepare_qat(
        self,
        graph: TraceGraph,
        is_input_quantized: typing.Optional[typing.Tuple[bool]] = None,
        backend: str = 'qnnpack',
        fuse_only: bool = False,
    ) -> torch.nn.Module:
        """Prepare model for QAT training

        Args:
            graph (TraceGraph): The computation graph of the model
            is_input_quantized (typing.Union[typing.Tuple[bool]], optional): Whether the input tensor(s) is (are) \
                quantized. Defaults to None.
            backend (str, optional): The backend of quantization. Defaults to 'qnnpack'.
            fuse_only (bool, optional): Whether the returned model is only fused in PostQuantizer. Defaults to False.

        Returns:
            torch.nn.Module: The QAT-ready model
        """

        graph.module.eval()

        self.prepare_qat_prep(graph, is_input_quantized, backend)
        if fuse_only:
            return graph.module

        # Unfornately, the suggested way below will try to fuse all the modules
        # even if some of the nodes are not in a quantized computation graph.
        # So we wrote some alternatives for the function.
        #   torch.quantization.prepare(graph.module, inplace=True)

        if hasattr(torch_q, 'get_default_qconfig_propagation_list'):
            whitelist = torch_q.get_default_qconfig_propagation_list()
        elif hasattr(torch_q, 'get_qconfig_propagation_list'):
            whitelist = torch_q.get_qconfig_propagation_list()
        else:
            whitelist = torch_q.DEFAULT_QCONFIG_PROPAGATE_WHITE_LIST

        if LooseVersion(torch.__version__) < LooseVersion("1.7.0"):
            torch_q.prepare(graph.module, inplace=True)

            if hasattr(graph.module, 'activation_post_process'):
                if hasattr(graph.module, "_forward_hooks"):
                    if len(graph.module._forward_hooks) > 0:
                        graph.module._forward_hooks.popitem()
                delattr(graph.module, "activation_post_process")
        else:
            torch_q.propagate_qconfig_(graph.module, qconfig_dict=None)
            for n in graph.forward_nodes:
                if not n.quantized:
                    if hasattr(n.module, "qconfig"):
                        delattr(n.module, "qconfig")

            if hasattr(torch_q, 'add_observer_'):
                add_observer_func = torch_q.add_observer_
            else:
                add_observer_func = sys.modules['torch.ao.quantization.quantize']._add_observer_

            if LooseVersion(torch.__version__) >= LooseVersion("1.8.0"):
                if LooseVersion(torch.__version__) >= LooseVersion("1.13.0"):
                    prepare_custom_config_dict = torch.ao.quantization.get_default_custom_config_dict()
                else:
                    prepare_custom_config_dict = {}

                custom_module_class_mapping = prepare_custom_config_dict.get(
                    "float_to_observed_custom_module_class", {}
                )

                custom_module_class_mapping.update(FUSE_PTQ_MODULES_CUSTOM)

                def patch_observer_set(orig_func):
                    def new_no_observer_set():
                        return set(FUSE_PTQ_MODULES_CUSTOM.values()) | orig_func()

                    return new_no_observer_set

                if LooseVersion(torch.__version__) >= LooseVersion("1.10.0"):
                    orig_no_observer_set = sys.modules['torch.ao.quantization.quantize'].no_observer_set
                    sys.modules['torch.ao.quantization.quantize'].no_observer_set = patch_observer_set(
                        orig_no_observer_set
                    )
                elif LooseVersion(torch.__version__) >= LooseVersion("1.9.0"):
                    orig_no_observer_set = torch.quantization.quantization_mappings.no_observer_set
                    torch.quantization.quantization_mappings.no_observer_set = patch_observer_set(orig_no_observer_set)

                add_observer_func(
                    graph.module,
                    qconfig_propagation_list=whitelist,
                    custom_module_class_mapping=custom_module_class_mapping,
                )
            else:
                add_observer_func(
                    graph.module,
                    qconfig_propagation_list=whitelist,
                )

        if self.dynamic_lstm_quant:
            mapping = {nn.LSTM: nnqd.LSTM}
            torch_q.convert(graph.module, mapping=mapping, inplace=True, remove_qconfig=False)

        for n in graph.forward_nodes:
            if not n.quantized:
                if hasattr(n.module, "_forward_hooks"):
                    if len(n.module._forward_hooks) > 0:
                        n.module._forward_hooks.popitem()
                if hasattr(n.module, "qconfig"):
                    delattr(n.module, "qconfig")
                if hasattr(n.module, "activation_post_process"):
                    delattr(n.module, "activation_post_process")

        if self.disable_requantization_for_cat:
            self.disable_requantization_for_cat_pass(graph)

        if self.quantized_input_stats is not None:
            self.prepare_quantized_inputs_pass(graph)

        if self.set_quantizable_op_stats:
            if self.quantized_op_stats is None:
                self.quantized_op_stats = {}
            self.quantized_op_stats.update(KNOWN_QSTATS)

        if self.quantized_op_stats is not None:
            self.prepare_quantized_ops_pass(graph)

        if self.quantize_op_action.get(nn.LSTM, None) and self.backend == 'qnnpack':
            self.fake_quantize_lstm_weights()

        return graph.module


class DynamicQuantizer(QATQuantizer):
    rewrite_graph: bool
    force_overwrite: bool
    is_input_quantized: typing.Optional[typing.Tuple[bool]]
    backend: str
    remove_weights_after_load: bool

    def __init__(self, model, dummy_input, work_dir: typing.Optional[str] = None, config: typing.Optional[dict] = None):
        """ Constructs a new DynamicQuantizer object

        Args:
            model: The model to be quantized
            dummy_input: A viable input to the model
            work_dir (typing.Optional[str], optional): The working directory in which the intermediate files will be \
                generated. Defaults to None, in which case "output" will be used.
            config (typing.Optional[dict]): Options for the quantizer
        """

        super().__init__(model, dummy_input, work_dir, config)

    def parse_config(self, config: dict):
        super().parse_config(config)

        self.rewrite_graph = False

        assert not self.asymmetric, "Asymmetric quantization is not supported for DynamicQuantizer"
        assert self.per_tensor, "Per-channel quantization is not supported for DynamicQuantizer"

    def quantize(self) -> nn.Module:
        """Prepare model for dynamic quantization"""

        self.model.eval()

        torch_q.quantize_dynamic(self.model, inplace=True)

        return self.model


class DeQuantizer(object):
    rewrite_graph: bool
    force_overwrite: bool
    remove_weights_after_load: bool
    extra_tracer_opts: typing.Optional[typing.Dict]
    target: str

    def __init__(self, model, dummy_input, work_dir: typing.Optional[str] = None, config: typing.Optional[dict] = None):
        """ Constructs a new DeQuantizer object

        Args:
            model: The model to be dequantized
            dummy_input: A viable input to the model
            work_dir (typing.Optional[str], optional): The working directory in which the intermediate files will be \
                generated. Defaults to None, in which case "output" will be used.
            config (typing.Optional[dict]): Options for the dequantizer
        """

        if isinstance(model, DataParallel) or isinstance(model, DistributedDataParallel):
            self.model = model.module
        else:
            self.model = model

        self.dummy_input = dummy_input
        self.work_dir = 'out' if work_dir is None else work_dir

        self.parse_config(config)

    def dequantize(self) -> nn.Module:
        # We need a model in training mode so that QAT could take place
        if self.target == 'qat':
            self.model.train()
        else:
            self.model.eval()

        # After tracing the model, we will get a TraceGraph object
        graph = trace(self.model, self.dummy_input, **self.extra_tracer_opts)

        if self.rewrite_graph:
            # Retrives the name of the model from type info
            model_name = type(self.model).__name__
            if model_name.endswith('_qat'):
                model_name = model_name[:-4]
            model_name_float = f'{model_name}_float'
            model_name_float_lower = model_name_float.lower()
            model_ns = f'tinynn_rewritten_models.{model_name_float}'

            model_code_path = os.path.join(self.work_dir, f'{model_name_float_lower}.py')
            model_weights_path = os.path.join(self.work_dir, f'{model_name_float_lower}.pth')

            # Generate the code for the modified model
            # We will try to do some op rewriting in the `rewrite_quantize_graph` function.
            if self.force_overwrite or not os.path.exists(model_code_path) or not os.path.exists(model_weights_path):
                self.rewrite_quantize_graph(graph)
                graph.generate_code(model_code_path, model_weights_path, model_name_float)

            # Import the new model
            rewritten_model = import_from_path(model_ns, model_code_path, model_name_float)()
            rewritten_model.load_state_dict(torch.load(model_weights_path))

            device = get_module_device(self.model)
            if device is not None:
                rewritten_model.to(device=device)

            # Remove the weights file to save space
            if self.remove_weights_after_load:
                os.unlink(model_weights_path)

            # Set the model to training mode before tracing again, since we need to perform QAT
            if self.target == 'qat':
                rewritten_model.train()
            else:
                rewritten_model.eval()

            # Update the model so that the original one can be released
            self.model = rewritten_model
        else:
            rewritten_model = self.model

        return rewritten_model

    def parse_config(self, config: typing.Optional[dict]):
        default_values = {
            'rewrite_graph': True,
            'force_overwrite': True,
            'extra_tracer_opts': {},
            'remove_weights_after_load': False,
            'target': 'qat',
        }

        if config is None:
            config = dict()

        for k, v in default_values.items():
            actual_v = config.get(k, v)
            setattr(self, k, actual_v)

    def rewrite_quantize_graph(self, graph: TraceGraph) -> None:
        """Rewrites the computation graph for quantization"""
        if graph.quantized:
            return

        graph_quantized = False
        for n in graph.forward_nodes:
            if n.type() in (torch_q.QuantStub, torch_q.DeQuantStub):
                graph_quantized = True
                break

        for n in graph.other_init_nodes:
            if n.type() is nnq.FloatFunctional:
                graph_quantized = True
                break

        if not graph_quantized:
            log.warning(
                'Graph is not quantized. No need to rewrite. Please pass in `config={"rewrite_graph": False}` to'
                ' suppress this warning'
            )
            return

        # Remove QuantStubs and DeQuantStubs
        names = [n.unique_name for n in graph.forward_nodes]
        for name in names:
            n = graph.nodes_map[name]
            if n.type() in (torch_q.QuantStub, torch_q.DeQuantStub):
                graph.remove_node(n)

        def _is_add_relu_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return (
                    cur_module.kind == 'add_relu'
                    and len(node.prev_nodes) > 1
                    and node.prev_nodes[0].type() is nnq.FloatFunctional
                )

        # Split FloatFunctional.add_relu to FloatFunctional.add and torch.relu
        add_relu_nodes = graph.filter_forward_nodes(_is_add_relu_node)
        for n in add_relu_nodes:
            n.module.kind = 'add'
            n.module.func_type = 'add'

            parts = n.module.full_name.split('.')[:-1] + [n.module.func_type]
            n.module.full_name = '.'.join(parts)

            with override_current_trace_graph(graph):
                next_func = TraceFunction('torch.relu').parse_args(n.next_tensors[0])

            next_out = torch.relu(n.next_tensors[0])
            graph.insert_after(n, next_func, [next_out])

        # Replace FloatFunctional.{add_scalar, mul_scalar} with torch.{add, mul}
        def _is_add_mul_scalar_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return (
                    cur_module.kind in ('add_scalar', 'mul_scalar')
                    and len(node.prev_nodes) > 1
                    and node.prev_nodes[0].type() is nnq.FloatFunctional
                )

        add_mul_scalar_nodes = graph.filter_forward_nodes(_is_add_mul_scalar_node)
        for n in add_mul_scalar_nodes:
            n.module.kind = n.module.kind.split('_')[0]
            n.module.func_type = n.module.kind

            parts = n.module.full_name.split('.')[:-1] + [n.module.func_type]
            n.module.full_name = '.'.join(parts)

        # Replace FloatFunctional.{add, mul, cat} with torch.{add, mul, cat}
        def _is_add_mul_cat_node(node: TraceNode, custom_data):
            cur_module = node.module
            cur_class = type(cur_module)
            if cur_class == TraceFunction:
                return (
                    cur_module.kind in ('add', 'mul', 'cat')
                    and len(node.prev_nodes) > 1
                    and node.prev_nodes[0].type() is nnq.FloatFunctional
                )

        add_mul_cat_nodes = graph.filter_forward_nodes(_is_add_mul_cat_node)
        for n in add_mul_cat_nodes:
            parts = ['torch'] + [n.module.func_type]
            n.module.full_name = '.'.join(parts)
            n.module.is_class = False

            n.prev_nodes[0].next_nodes.remove(n)
            n.prev_nodes.pop(0)
            n.prev_indices.pop(0)

            n.module.tensor_names.pop(0)
            n.module.args_parsed.pop(0)
            n.module.args_to_string(n.module.args_parsed)

        # Remove FloatFunctional nodes
        names = [n.unique_name for n in graph.other_init_nodes]
        for name in names:
            n = graph.nodes_map[name]
            if n.type() is nnq.FloatFunctional:
                graph.other_init_nodes.remove(n)
                del graph.nodes_map[n.unique_name]

        graph.quantized = False
        graph.recompute_forward_order()


def load_creation_func_names():
    if len(creation_func_names) == 0:
        funcs_d = load_creation_funcs()
        for ns, funcs_v in funcs_d.items():
            ns_str = qualified_name(ns)
            creation_func_names.extend([f'{ns_str}.{x}' for x in funcs_v])
    return creation_func_names


def get_dict_from_rules(rules):
    rule_dict = {}
    for rule in rules:
        base_rule_dict = rule_dict
        for module_cls in reversed(rule):
            # Node properties (has_key, child_nodes)
            base_rule_dict.setdefault(module_cls, [False, {}])
            base_rule_pair = base_rule_dict[module_cls]
            base_rule_dict = base_rule_pair[1]
        base_rule_pair[0] = True
    return rule_dict


def load_processed_qat_rules():
    if len(processed_qat_rules) == 0:
        # Constructor a prefix tree for the QAT rules
        fuse_rules = sorted(FUSE_RULE_LIST, key=lambda x: len(x), reverse=True)
        rule_dict = get_dict_from_rules(fuse_rules)
        processed_qat_rules.update(rule_dict)
    return processed_qat_rules


def load_processed_ptq_rules():
    if len(processed_ptq_rules) == 0:
        # Constructor a prefix tree for the PTQ rules
        filtered_ptq_rules = {
            k
            for k, v in FUSE_RULE_LIST_PTQ_ONLY.items()
            if v is None or LooseVersion(torch.__version__) >= LooseVersion(v)
        }
        ptq_rules = set(FUSE_RULE_LIST).union(set(filtered_ptq_rules))
        fuse_rules = sorted(ptq_rules, key=lambda x: len(x), reverse=True)
        rule_dict = get_dict_from_rules(fuse_rules)
        processed_ptq_rules.update(rule_dict)
    return processed_ptq_rules


def load_processed_extra_q_rules():
    if len(processed_extra_q_rules) == 0:
        # Constructor a prefix tree for the QAT rules
        fuse_rules = sorted(FUSE_RULE_LIST_EXTRA, key=lambda x: len(x), reverse=True)
        rule_dict = get_dict_from_rules(fuse_rules)
        processed_extra_q_rules.update(rule_dict)
    return processed_extra_q_rules


def load_processed_rewrite_to_fuse_rules():
    if len(processed_rewrite_to_fuse_rules) == 0:
        # Constructor a prefix tree for the rewrite rules
        fuse_rules = sorted(REWRITE_TO_FUSE_RULE_LIST, key=lambda x: len(x), reverse=True)
        rule_dict = get_dict_from_rules(fuse_rules)
        processed_rewrite_to_fuse_rules.update(rule_dict)
    return processed_rewrite_to_fuse_rules


def load_processed_all_qat_rules():
    if len(processed_all_qat_rules) == 0:
        # Constructor a prefix tree for the QAT rules
        qat_rules = set(FUSE_RULE_LIST).union(set(FUSE_RULE_LIST_EXTRA))
        fuse_rules = sorted(qat_rules, key=lambda x: len(x), reverse=True)
        rule_dict = get_dict_from_rules(fuse_rules)
        processed_all_qat_rules.update(rule_dict)
    return processed_all_qat_rules


def load_processed_all_ptq_rules():
    if len(processed_all_ptq_rules) == 0:
        # Constructor a prefix tree for the QAT rules
        filtered_ptq_rules = {
            k
            for k, v in FUSE_RULE_LIST_PTQ_ONLY.items()
            if v is None or LooseVersion(torch.__version__) >= LooseVersion(v)
        }
        ptq_rules = set(FUSE_RULE_LIST).union(set(filtered_ptq_rules)).union(set(FUSE_RULE_LIST_EXTRA))
        fuse_rules = sorted(ptq_rules, key=lambda x: len(x), reverse=True)
        rule_dict = get_dict_from_rules(fuse_rules)
        processed_all_ptq_rules.update(rule_dict)
    return processed_all_ptq_rules