# Owner(s): ["oncall: quantization"]

import collections
import copy
import math
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.intrinsic as nni
import torch.nn.quantized as nnq
toq = torch.ops.quantized
from torch.testing._internal.common_quantization import (
    skipIfNoFBGEMM,
    skip_if_no_torchvision,
    QuantizationTestCase,
    NodeSpec,
)
from torch.testing import FileCheck
from torch.quantization import (
    ObserverBase,
    FakeQuantizeBase,
    QConfig,
    MinMaxObserver,
)
from torch.quantization.quantize_fx import (
    prepare_fx,
    convert_fx,
)
from torch.ao.quantization._dbr.quantization_state import AutoQuantizationState

import torch.ao.quantization._quantize_dbr as _quantize_dbr
import torch.ao.ns._numeric_suite_dbr as ns
# TODO(future PR): move these utils out of the FX folder
import torch.ao.ns._numeric_suite_fx as ns_fx
from torch.ao.quantization._dbr.torchscript_utils import (
    remove_redundant_aliases,
)

def _allclose(a, b):
    if isinstance(a, tuple):
        assert isinstance(b, tuple)
        result = True
        for a_inner, b_inner in zip(a, b):
            result = result and torch.allclose(a_inner, b_inner)
        return result
    elif isinstance(a, torch.Tensor):
        assert isinstance(b, torch.Tensor)
        return torch.allclose(a, b)
    raise AssertionError('unhandled type')


class QuantizeDBRTestCase(QuantizationTestCase):
    def _test_auto_tracing(
        self,
        m,
        qconfig,
        example_args,
        fuse_modules=True,
        do_fx_comparison=True,
        do_torchscript_checks=True,
        # there are some keys in DBR prepare_custom_config_dict which
        # are not supported in FX, this argument is for DBR only
        dbr_prepare_custom_config_dict=None,
    ):
        m_copy = copy.deepcopy(m)

        qconfig_dict = {'': qconfig}

        mp = _quantize_dbr.prepare(
            m, qconfig_dict, example_args, fuse_modules=fuse_modules,
            prepare_custom_config_dict=dbr_prepare_custom_config_dict)
        out_p = mp(*example_args)
        # print(mp)
        mq = _quantize_dbr.convert(mp)
        # print(mq)
        # verify it runs
        out_q = mq(*example_args)
        # print(out_q)

        # compare it against FX
        if do_fx_comparison:
            m_copy_p = prepare_fx(m_copy, {'': qconfig})
            out_m_copy_p = m_copy_p(*example_args)
            # print(m_copy_p)
            m_copy_q = convert_fx(m_copy_p)
            # print(m_copy_q)
            # print(m_copy_q.graph)
            out_q_fx = m_copy_q(*example_args)
            # print(out_q)
            # print(out_q_fx)
            self.assertTrue(_allclose(out_p, out_m_copy_p))
            # print(out_q)
            # print(out_q_fx)
            self.assertTrue(_allclose(out_q, out_q_fx))

        if do_torchscript_checks:
            # verify torch.jit.trace works
            mq_jit_traced = torch.jit.trace(
                mq, example_args, check_trace=False)
            # print(mq_jit_traced.graph)
            traced_out = mq_jit_traced(*example_args)
            self.assertTrue(_allclose(traced_out, out_q))

            # verify torch.jit.script works
            rewritten = mq.rewrite_for_scripting()
            rewritten_out = rewritten(*example_args)
            # print(rewritten)
            self.assertTrue(_allclose(rewritten_out, out_q))

            scripted_rewritten = torch.jit.script(rewritten)
            # print(scripted_rewritten.graph)
            scripted_rewritten_out = scripted_rewritten(*example_args)
            # print('scripted_rewritten_out', scripted_rewritten_out)
            self.assertTrue(_allclose(scripted_rewritten_out, out_q))

            traced_rewritten = torch.jit.trace(
                rewritten, example_args, check_trace=False)
            traced_rewritten_out = traced_rewritten(*example_args)
            self.assertTrue(_allclose(traced_rewritten_out, out_q))

@skipIfNoFBGEMM
class TestQuantizeDBRIndividualOps(QuantizeDBRTestCase):
    """
    Tests that DBR quantization covers individual ops
    """
    def test_conv(self):
        convs = {1: nn.Conv1d, 2: nn.Conv2d, 3: nn.Conv3d}

        class M(torch.nn.Module):
            def __init__(self, dim):
                super().__init__()
                self.conv = convs[dim](3, 3, 3)

            def forward(self, x):
                x1 = self.conv(x)
                return x1

        data = {
            1: torch.randn(1, 3, 10),
            2: torch.randn(1, 3, 10, 10),
            3: torch.randn(1, 3, 5, 5, 5)
        }
        for dim in range(1, 4):
            m = M(dim).eval()
            qconfig = torch.quantization.default_qconfig
            self._test_auto_tracing(m, qconfig, (data[dim],))

    def test_conv_functional(self):
        convs = {1: F.conv1d, 2: F.conv2d, 3: F.conv3d}

        class M(torch.nn.Module):
            def __init__(self, dim, weight, bias):
                super().__init__()
                self.conv_func = convs[dim]
                self.weight = torch.nn.Parameter(weight)
                self.bias = torch.nn.Parameter(bias)
                self.stride = (1,) * dim
                self.padding = (0,) * dim
                self.dilation = (1,) * dim
                self.groups = 1

            def forward(self, x):
                x = self.conv_func(
                    x, self.weight, self.bias, self.stride, self.padding,
                    self.dilation, self.groups)
                return x

        data = {
            1: torch.randn(1, 3, 10),
            2: torch.randn(1, 3, 10, 10),
            3: torch.randn(1, 3, 5, 5, 5)
        }
        bias = torch.randn(1)
        for dim in range(1, 4):
            model_fp32 = M(dim, data[dim], bias).eval()
            qconfig = torch.quantization.default_qconfig
            self._test_auto_tracing(model_fp32, qconfig, (data[dim],))

    def test_linear_dynamic(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(1, 1)

            def forward(self, x):
                x1 = self.linear(x)
                return x1

        m = M().eval()
        qconfig = torch.quantization.default_dynamic_qconfig
        # qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 1, 1),))

    def test_linear_functional(self):
        class LinearFunctional(nn.Module):
            def __init__(self):
                super().__init__()
                self.w1 = nn.Parameter(torch.empty(4, 4))
                self.b1 = nn.Parameter(torch.ones(4))
                torch.nn.init.kaiming_uniform_(self.w1, a=math.sqrt(5))

            def forward(self, x):
                x = F.linear(x, self.w1, self.b1)
                return x

        model_fp32 = LinearFunctional().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.randn(1, 1, 4, 4),))

    def test_linear_functional_nobias(self):
        class LinearFunctional(nn.Module):
            def __init__(self):
                super().__init__()
                self.w1 = nn.Parameter(torch.empty(4, 4))
                torch.nn.init.kaiming_uniform_(self.w1, a=math.sqrt(5))

            def forward(self, x):
                x = F.linear(x, self.w1)
                return x

        model_fp32 = LinearFunctional().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.randn(1, 1, 4, 4),))

    # TODO(future PR): implement observer sharing to match FX
    def test_cat_fp32(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = torch.cat([x, x], dim=1)
                return x

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 2, 2),))

        class M(torch.nn.Module):
            def forward(self, x):
                x = torch.cat((x, x), dim=1)
                return x

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_cat_int(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = torch.cat([x, x], dim=1)
                return x

        qconfig = torch.quantization.default_qconfig
        for dtype in (torch.int32, torch.int64):
            m = M().eval()
            self._test_auto_tracing(
                m, qconfig, (torch.zeros(1, 1, 1, 1, dtype=dtype),),
                # FX graph mode quant does not support this yet
                do_fx_comparison=False)

    def test_add(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = x + x
                x = x + 1.0
                x = 1.0 + x
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_add_int32(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = x + x
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.ones(1, 1, 2, 2, dtype=torch.int32),),
            # FX graph mode quantization does not automatically detect
            # tensor inputs in non-float dtypes.
            do_fx_comparison=False)

    def test_sub(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = x - x
                x = x - 1.0
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_mul(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = x * x
                x = x * 1.0
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_mul_int(self):
        # TODO: make all the math functions work correctly for integer types
        # TODO: make the same improvement in FX graph mode quant, if possible
        class M(torch.nn.Module):
            def forward(self, x):
                x = x * x
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        for dtype in (torch.int32, torch.int64):
            self._test_auto_tracing(
                copy.deepcopy(model_fp32), qconfig,
                (torch.ones(1, 1, 2, 2, dtype=dtype),),
                # FX graph mode quant does not support this yet
                do_fx_comparison=False)

    def test_div(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = x / x
                x = x / 1.0
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_embedding(self):
        # Note: this test is just testing that models with embeddings
        # do not crash with a global qconfig defined. Embedding quantization
        # is not actually happening in this prototype yet.
        # TODO(future PR): fix this and update this code.

        # test subclass
        class EmbeddingSubclass(nn.Embedding):
            pass

        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.embedding = EmbeddingSubclass(1, 1)

            def forward(self, x):
                x = self.embedding(x)
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_dynamic_qconfig
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.LongTensor([[0]]),),
            fuse_modules=False)

        # test regular embedding
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.embedding = nn.Embedding(1, 1)

            def forward(self, x):
                x = self.embedding(x)
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_dynamic_qconfig
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.LongTensor([[0]]),),
            fuse_modules=False)


@skipIfNoFBGEMM
class TestQuantizeDBR(QuantizeDBRTestCase):
    def test_fusion(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)
                self.relu = torch.nn.ReLU()
                self.child = nn.Sequential(
                    nn.Conv2d(1, 1, 1),
                    nn.ReLU(),
                )

            def forward(self, x):
                x = self.conv(x)
                x = self.relu(x)
                x = self.child(x)
                return x

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        mp = _quantize_dbr.prepare(m, {'': qconfig}, (torch.randn(1, 1, 1, 1),))
        self.assertTrue(isinstance(mp.conv, nni.ConvReLU2d))
        self.assertTrue(isinstance(mp.child[0], nni.ConvReLU2d))

    def test_fusion2(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)
                self.bn = torch.nn.BatchNorm2d(1)
                # self.conv2 = torch.nn.Conv2d(1, 1, 1)
                self.relu = torch.nn.LeakyReLU()

            def forward(self, x):
                x = self.conv(x)
                x = self.bn(x)
                x = self.relu(x)
                return x

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_fusion_called_multiple_times(self):
        """
        Tests that fusion works if the modules to fuse get called multiple
        times in the same forward.

        Currently, observers are not shared between successive calls of
        the same module.
        TODO(future PR): make them shared (this is easy to detect)
        """
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)
                self.relu = torch.nn.ReLU()

            def forward(self, x):
                for _ in range(2):
                    x = self.conv(x)
                    x = self.relu(x)
                return x

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        # fx graph mode quant doesn't support using a single module multiple times
        # right now, so this would crash, we can handle this case later
        # if it is needed
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 2, 2),), do_fx_comparison=False)

    def test_fusion_functions(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = x + x
                x = torch.relu(x)
                return x

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        mp = _quantize_dbr.prepare(m, {'': qconfig}, (torch.randn(1, 1, 1, 1),))

        self.assertTrue(
            mp._auto_quant_state.idx_to_seen_q_op_infos[0].fusion_info is not None)
        self.assertTrue(
            mp._auto_quant_state.idx_to_seen_q_op_infos[1].fusion_info is not None)

        # verify that the add relu is not observed
        self.assertTrue(
            '1' not in mp._auto_quant_state.tensor_id_to_observer)
        # verify that the relu is observed
        self.assertTrue(
            '2' in mp._auto_quant_state.tensor_id_to_observer)

        mp(torch.randn(1, 1, 1, 1))
        mq = _quantize_dbr.convert(mp)

        # verify that the add-relu got fused
        mqt = torch.jit.trace(mq, (torch.randn(1, 1, 1, 1),))
        FileCheck().check_count("quantized::add_relu", 1, exactly=True).run(
            mqt.graph)

        # TODO(future PR): use information about non-quantizeable ops during
        #   matching fusion patterns

    def test_observers_not_touched_by_tracing(self):
        """
        Verifies that running dynamic tracing does not change any data
        stored in observers and fake quants.
        """
        m = nn.Sequential(nn.Conv2d(1, 1, 1)).eval()
        qconfig = torch.quantization.default_qconfig
        mp = _quantize_dbr.prepare(m, {'': qconfig}, (torch.randn(1, 1, 1, 1),))
        for _, mod in mp.named_modules():
            if isinstance(mod, (ObserverBase, FakeQuantizeBase)):
                scale, zp = mod.calculate_qparams()
                # Assume that if scale is 1.0 and zp is 0, no calibration
                # has happened.
                self.assertTrue(torch.allclose(scale, torch.ones(1)))
                self.assertTrue(torch.equal(zp, torch.zeros(1, dtype=torch.long)))

    def test_multiple_modules(self):
        m = nn.Sequential(
            nn.Sequential(nn.Conv2d(1, 1, 1)),
            nn.Sequential(nn.Conv2d(1, 1, 1)),
        ).eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_child_modules(self):
        m = nn.Sequential(nn.Sequential(nn.Conv2d(1, 1, 1))).eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_conv_mod_qat(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)

            def forward(self, x):
                x1 = self.conv(x)
                return x1

        m = M().eval()
        qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
        self._test_auto_tracing(
            copy.deepcopy(m), qconfig, (torch.randn(1, 1, 2, 2),))

        # test backprop does not crash
        inputs = torch.randn(1, 1, 1, 1)
        inputs.requires_grad = True
        mp = _quantize_dbr.prepare(m, {'': qconfig}, (inputs,))
        output = mp(inputs)
        labels = torch.randn(1, 1, 1, 1)
        loss = (output - labels).sum()
        loss.backward()
        optim = torch.optim.SGD(mp.parameters(), lr=0.01)
        optim.step()

    def test_conv_functional_qat(self):

        class M(torch.nn.Module):
            def __init__(self, weight2d, bias2d):
                super().__init__()
                self.weight2d = torch.nn.Parameter(weight2d)
                self.bias2d = torch.nn.Parameter(bias2d)
                self.stride2d = (1, 1)
                self.padding2d = (0, 0)
                self.dilation2d = (1, 1)
                self.groups = 1

            def forward(self, x):
                x = F.conv2d(
                    x, self.weight2d, self.bias2d, self.stride2d, self.padding2d,
                    self.dilation2d, self.groups)
                return x

        m = M(torch.randn(1, 1, 1, 1), torch.randn(1)).eval()
        qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 2, 2),))

        # test backprop does not crash
        inputs = torch.randn(1, 1, 1, 1)
        inputs.requires_grad = True
        m = M(torch.randn(1, 1, 1, 1), torch.randn(1)).eval()
        mp = _quantize_dbr.prepare(m, {'': qconfig}, (inputs,))
        output = mp(inputs)
        labels = torch.randn(1, 1, 1, 1)
        loss = (output - labels).sum()
        loss.backward()
        optim = torch.optim.SGD(mp.parameters(), lr=0.01)
        optim.step()

    @unittest.skip('FX graph mode is using fake_quantize with PTQ, TODO verify')
    def test_conv_unsupported_inplace_conv(self):
        """
        Verifies that having an quantizeable op which is inplace
        is handled well
        """
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)
                self.conv2 = torch.nn.Conv2d(1, 1, 1)

            def forward(self, x):
                x = self.conv(x)
                x = F.hardsigmoid(x, inplace=True)
                x = self.conv2(x)
                return x

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_conv_functional_dynamic_weights(self):
        class M(torch.nn.Module):
            def __init__(self, weight2d, bias2d):
                super().__init__()
                self.weight2d = torch.nn.Parameter(weight2d)
                self.bias2d = torch.nn.Parameter(bias2d)
                self.stride2d = (1, 1)
                self.padding2d = (0, 0)
                self.dilation2d = (1, 1)
                self.groups = 1

            def forward(self, x):
                updated_weight = self.weight2d * x
                x = F.conv2d(
                    x, updated_weight, self.bias2d, self.stride2d, self.padding2d,
                    self.dilation2d, self.groups)
                return x

        model_fp32 = M(torch.randn(1, 1, 1, 1), torch.randn(1)).eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.randn(1, 1, 2, 2),),
            # FX implements this functionality instead of skipping it
            do_fx_comparison=False,
            # TODO enable scripting support for this
            do_torchscript_checks=False)

    def test_method(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = x + x
                x = torch.relu(x)
                # x = x.relu()
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_module_created_during_forward(self):
        """Some BERT models have this pattern"""
        class M(torch.nn.Module):
            def forward(self, x):
                x = nn.Softmax(dim=-1)(x)
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.randn(1, 1, 1, 1),),
            # This syntax is not supported by FX or TorchScript
            do_fx_comparison=False, do_torchscript_checks=False)

    def test_module_returns_namedtuple(self):
        NamedTuple = collections.namedtuple("NamedTuple", ["x0", "x1"])

        """Some hf models have this pattern"""
        class M1(torch.nn.Module):
            def forward(self, x):
                return NamedTuple(x, x)

        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.m1 = M1()

            def forward(self, x):
                m1 = self.m1(x)
                return (m1.x0, m1.x1)

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.randn(1, 1, 1, 1),),
            # TODO(future PR): add FX rewrite support
            do_fx_comparison=False, do_torchscript_checks=False)

    def test_child_module_does_not_return_tensor(self):
        class M1(torch.nn.Module):
            def forward(self, x):
                pass

        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.m1 = M1()

            def forward(self, x):
                self.m1(x)
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.randn(1, 1, 1, 1),),
            # TODO(future PR): add FX rewrite support
            do_fx_comparison=False, do_torchscript_checks=False)

    def _get_non_traceable_module_class_test_model(self):
        class M1(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)

            def forward(self, x):
                x = self.conv(x)
                x = x + x
                return x

        class M2(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.m1 = M1()
                self.conv = torch.nn.Conv2d(1, 1, 1)

            def forward(self, x):
                x = self.m1(x)
                x = self.conv(x)
                x = x + x
                return x

        class M3(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.m2 = M2()

            def forward(self, x):
                x = self.m2(x)
                return x

        return M3().eval(), M1, M2, M3

    def test_prepare_custom_config_dict_non_traceable_module_class_child_leaf(self):

        # if M1 is set as leaf, M2 and M3 should have auto_quant_state
        qconfig_dict = {'': torch.quantization.default_qconfig}
        m, M1, M2, M3 = self._get_non_traceable_module_class_test_model()
        prepare_custom_config_dict = {
            'non_traceable_module_class': [M1],
        }
        mp = _quantize_dbr.prepare(
            m, qconfig_dict, (torch.randn(1, 1, 1, 1),),
            prepare_custom_config_dict=prepare_custom_config_dict)
        self.assertTrue(not hasattr(mp.m2.m1, '_auto_quant_state'))
        self.assertTrue(hasattr(mp.m2, '_auto_quant_state'))
        self.assertTrue(hasattr(mp, '_auto_quant_state'))
        mq = _quantize_dbr.convert(mp)
        self.assertTrue(isinstance(mq.m2.m1.conv, nn.Conv2d))
        self.assertTrue(isinstance(mq.m2.conv, nnq.Conv2d))
        mqt = torch.jit.trace(mq, (torch.randn(1, 1, 1, 1),))

        # mqt.m2.m1 should not have quantized ops
        FileCheck().check_count("aten::add", 1, exactly=True).run(mqt.m2.m1.graph)
        FileCheck().check_count("quantized::add", 0, exactly=True).run(mqt.m2.m1.graph)
        # mqt.m2.m1 should have quantized ops
        FileCheck().check_count("aten::add", 0, exactly=True).run(mqt.m2.graph)
        FileCheck().check_count("quantized::add", 1, exactly=True).run(mqt.m2.graph)

        # TODO(future PR): ensure modules in leaves do not get quantized

    def test_prepare_custom_config_dict_non_traceable_module_class_mid_leaf(self):
        # if M2 is set as leaf, only M1 should have auto_quant_state
        qconfig_dict = {'': torch.quantization.default_qconfig}
        m, M1, M2, M3 = self._get_non_traceable_module_class_test_model()
        prepare_custom_config_dict = {
            'non_traceable_module_class': [M2],
        }
        mp = _quantize_dbr.prepare(
            m, qconfig_dict, (torch.randn(1, 1, 1, 1),),
            prepare_custom_config_dict=prepare_custom_config_dict)
        self.assertTrue(not hasattr(mp.m2.m1, '_auto_quant_state'))
        self.assertTrue(not hasattr(mp.m2, '_auto_quant_state'))
        self.assertTrue(hasattr(mp, '_auto_quant_state'))
        mq = _quantize_dbr.convert(mp)
        self.assertTrue(isinstance(mq.m2.m1.conv, nn.Conv2d))
        self.assertTrue(isinstance(mq.m2.conv, nn.Conv2d))
        mqt = torch.jit.trace(mq, (torch.randn(1, 1, 1, 1),))

        # mqt.m2 and all children should not have quantized ops
        FileCheck().check_count("aten::add", 1, exactly=True).run(mqt.m2.m1.graph)
        FileCheck().check_count("quantized::add", 0, exactly=True).run(mqt.m2.m1.graph)
        FileCheck().check_count("aten::add", 1, exactly=True).run(mqt.m2.graph)
        FileCheck().check_count("quantized::add", 0, exactly=True).run(mqt.m2.graph)

    def test_module_list(self):
        class Child(torch.nn.Module):
            def forward(self, x):
                return x + x

        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.module_list = torch.nn.ModuleList([
                    Child(),
                ])

            def forward(self, x):
                for module in self.module_list:
                    # TODO(future PR): we should see if there is a better
                    # solution other than asking users to do this
                    if not isinstance(module, AutoQuantizationState):
                        x = module(x)
                return x

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            m, qconfig, (torch.randn(8, 1, 1, 1),),
            # TODO(future PR): enable scripting for ModuleList + DBR
            do_fx_comparison=True, do_torchscript_checks=False)

    @unittest.skip('TODO build this')
    def test_module_input_types(self):
        class M(torch.nn.Module):
            def forward(self, x=None, y=None):
                print('x', x)
                print('y', y)
                assert x is not None and y is not None
                return (x, y)

        model_fp32 = M().eval()
        example_inputs = {'y': torch.randn(1), 'x': torch.randn(1)}
        ExampleInputsTupleCtr = collections.namedtuple('ExampleInputs', example_inputs)
        example_inputs_tuple = ExampleInputsTupleCtr(**example_inputs)
        ms = torch.jit.trace(model_fp32, example_inputs_tuple)

        return
        qconfig = torch.quantization.default_qconfig

        # dict
        kwargs = {'x': torch.randn(1, 1, 2, 2)}
        self._test_auto_tracing(model_fp32, qconfig, (), kwargs)

    def test_module_return_types(self):
        class M1(torch.nn.Module):
            def forward(self, x):
                return x, x

        class M2(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.m1 = M1()

            def forward(self, x):
                x1, x2 = self.m1(x)
                return x1

        model_fp32 = M2().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_inplace_unquantizeable_op(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv1 = nn.Conv2d(1, 1, 1)
                self.silu = nn.SiLU(inplace=True)
                # self.silu = nn.SiLU()
                self.conv2 = nn.Conv2d(1, 1, 1)

            def forward(self, x):
                x = self.conv1(x)
                x = self.silu(x)
                x = self.conv2(x)
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_vovnet_sequential(self):
        # We cannot quantize SequentialAppendList directly because
        # AutoQuantizationStateModuleDict would appear in self.items.
        # However, we can wrap it and quantize the wrapper.
        class SequentialAppendList(nn.Sequential):
            def __init__(self, *args):
                super(SequentialAppendList, self).__init__(*args)

            def forward(self, x: torch.Tensor) -> torch.Tensor:
                concat_list = []
                for i, module in enumerate(self):
                    if i == 0:
                        concat_list.append(module(x))
                    else:
                        concat_list.append(module(concat_list[-1]))
                x = torch.cat(concat_list, dim=1)
                return x

        class Wrapper(nn.Module):
            def __init__(self, *args):
                super().__init__()
                self.append_list = SequentialAppendList(*args)

            def forward(self, x):
                x = self.append_list(x)
                return x

        m = Wrapper(torch.nn.Conv2d(1, 1, 1)).eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 1, 1),))

    def test_unsupported_ops(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = F.tanhshrink(x)
                x = x + x
                x = F.tanhshrink(x)
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_unsupported_ops_recorded(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv2d = nn.Conv2d(1, 1, 1)
                self.softshrink = nn.Softshrink()

            def forward(self, x):
                # supported
                x = self.conv2d(x)
                x = x + x
                # not supported
                x = self.softshrink(x)
                x = F.tanhshrink(x)
                return x

        m = M().eval()
        qconfig_dict = {'': torch.quantization.default_qconfig}
        mp = _quantize_dbr.prepare(m, qconfig_dict, (torch.randn(1, 1, 1, 1),))
        self.assertTrue(len(mp._auto_quant_state.seen_nonq_op_infos) == 2)
        self.assertTrue(mp._auto_quant_state.seen_nonq_op_infos[0].type == nn.Softshrink)
        self.assertTrue(mp._auto_quant_state.seen_nonq_op_infos[1].type == F.tanhshrink)

    def test_unknown_op_after_quantized(self):
        class M(torch.nn.Module):
            def forward(self, x):
                x = x + x
                std = x.std()
                return std

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.randn(1, 1, 2, 2),),
            fuse_modules=False)

    def test_module_calls_items(self):
        # We cannot quantize M1 directly because
        # AutoQuantizationStateModuleDict would appear in self.items.
        # However, we can wrap it and quantize the wrapper.
        class M1(torch.nn.ModuleDict):
            def __init__(self):
                super().__init__()
                for i in range(2):
                    layer = nn.ReLU()
                    self.add_module("layer_" + str(i), layer)

            def forward(self, x):
                layers = [x]
                for name, layer in self.items():
                    layers.append(layer(x))
                return torch.cat(layers, dim=1)

        class M2(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.m1 = M1()

            def forward(self, x):
                x = self.m1(x)
                return x

        model_fp32 = M2().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.randn(1, 1, 2, 2),),
            # TODO(future PR): implement observer sharing for torch.cat
            # in DBR quant, to ensure that numerical behavior matches
            do_fx_comparison=False)

    def test_subclass_of_quantizeable_module(self):
        """
        If a user creates a subclass of nn.BatchNorm2d, that subclass
        should not be quantized unless the user defines a custom module.
        """
        class BN2d(torch.nn.BatchNorm2d):
            pass

        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)
                self.bn = BN2d(1)
                self.conv2 = torch.nn.Conv2d(1, 1, 1)

            def forward(self, x):
                x = self.conv(x)
                x = self.bn(x)
                x = self.conv2(x)
                return x

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            m, qconfig, (torch.randn(1, 1, 2, 2),),
            # the module is not symbolically traceable
            do_fx_comparison=False)

    # TODO(future PR): move into a separate test file
    def test_numeric_suite(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = nn.Conv2d(1, 1, 1)
                self.conv2 = nn.Sequential(nn.Conv2d(1, 1, 1))

            def forward(self, x):
                x = self.conv(x)
                x = self.conv2(x)
                x = x + x
                return x

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        example_args = (torch.randn(1, 1, 2, 2),)
        mp = _quantize_dbr.prepare(m, {'': qconfig}, example_args)
        out_p = mp(*example_args)
        mq = _quantize_dbr.convert(copy.deepcopy(mp))
        out_q = mq(*example_args)

        mp, mq = ns.add_loggers('mp', mp, 'mq', mq)

        mp(*example_args)
        mq(*example_args)

        act_comparison = ns.extract_logger_info(mp, mq, 'mq')
        ns_fx.extend_logger_results_with_comparison(
            act_comparison, 'mp', 'mq', torch.ao.ns.fx.utils.compute_sqnr,
            'sqnr')

        # TODO(future PR): enforce validity of the result above, using
        # NS for FX utils. Will need some refactoring.

        # TODO(future PR): consider adding a util for below
        to_print = []
        for idx, (layer_name, v) in enumerate(act_comparison.items()):
            to_print.append([
                layer_name,
                v['node_output']['mq'][0]['fqn'],
                v['node_output']['mq'][0]['ref_node_target_type'],
                v['node_output']['mq'][0]['sqnr']])

    def test_qconfig_dict_global(self):
        """
        Verifies that the '' option of qconfig_dict works
        """

        # regular case
        m = nn.Sequential(nn.Conv2d(1, 1, 1))
        qconfig_dict = {'': torch.quantization.default_qconfig}
        example_args = (torch.randn(1, 1, 1, 1),)
        mp = _quantize_dbr.prepare(m, qconfig_dict, example_args)
        mp(*example_args)
        mq = _quantize_dbr.convert(mp)
        mq(*example_args)
        self.assertTrue(isinstance(mq[0], nnq.Conv2d))

        # quantization turned off
        m = nn.Sequential(nn.Conv2d(1, 1, 1))
        qconfig_dict = {'': None}
        example_args = (torch.randn(1, 1, 1, 1),)
        mp = _quantize_dbr.prepare(m, qconfig_dict, example_args)
        mp(*example_args)
        mq = _quantize_dbr.convert(mp)
        mq(*example_args)
        self.assertTrue(isinstance(mq[0], nn.Conv2d))

    def test_qconfig_dict_object_type_module(self):
        """
        Verifies that the 'object_type' option of qconfig_dict works
        on module types.
        """
        m = nn.Sequential(
            nn.Conv2d(1, 1, 1),
            nn.Hardswish(),
            nn.Conv2d(1, 1, 1),
        )
        qconfig_dict = {
            '': torch.quantization.default_qconfig,
            'object_type': [
                (nn.Conv2d, torch.quantization.default_qconfig),
                (nn.Hardswish, None),
            ],
        }
        example_args = (torch.randn(1, 1, 1, 1),)
        mp = _quantize_dbr.prepare(m, qconfig_dict, example_args)
        mp(*example_args)
        mq = _quantize_dbr.convert(mp)
        mq(*example_args)
        self.assertTrue(isinstance(mq[0], nnq.Conv2d))
        self.assertTrue(isinstance(mq[1], nn.Hardswish))
        self.assertTrue(isinstance(mq[2], nnq.Conv2d))

    def test_qconfig_dict_object_type_function(self):
        """
        Verifies that the 'object_type' option of qconfig_dict works
        on function types.
        """
        class M(nn.Module):
            def forward(self, x):
                x = x + x
                x = x * x
                return x

        m = M()
        qconfig_dict = {
            '': torch.quantization.default_qconfig,
            'object_type': [
                (torch.add, None),
            ],
        }
        example_args = (torch.randn(1, 1, 1, 1),)
        mp = _quantize_dbr.prepare(m, qconfig_dict, example_args)
        mp(*example_args)
        mq = _quantize_dbr.convert(mp)
        mq(*example_args)
        rewritten = mq.rewrite_for_scripting()
        expected_occurrence = {
            NodeSpec.call_function(torch.add): 1,
            NodeSpec.call_function(toq.add): 0,
            NodeSpec.call_function(toq.mul): 1,
        }
        self.checkGraphModuleNodes(
            rewritten, expected_node_occurrence=expected_occurrence)

    def test_qconfig_dict_object_type_method(self):
        """
        Verifies that the 'object_type' option of qconfig_dict works
        on method types.
        """
        class M(nn.Module):
            def forward(self, x):
                x = x.add(x)
                x = x.mul(x)
                return x

        m = M()
        qconfig_dict = {
            '': torch.quantization.default_qconfig,
            'object_type': [
                (torch.Tensor.add, None),
            ],
        }
        example_args = (torch.randn(1, 1, 1, 1),)
        mp = _quantize_dbr.prepare(m, qconfig_dict, example_args)
        mp(*example_args)
        mq = _quantize_dbr.convert(mp)
        mq(*example_args)
        rewritten = mq.rewrite_for_scripting()
        expected_occurrence = {
            NodeSpec.call_function(torch.add): 1,
            NodeSpec.call_function(toq.add): 0,
            NodeSpec.call_function(toq.mul): 1,
        }
        self.checkGraphModuleNodes(
            rewritten, expected_node_occurrence=expected_occurrence)

    def test_qconfig_dict_object_type_function_global_none(self):
        """
        Verifies that the 'object_type' option of qconfig_dict works
        on function types when global qconfig is None.
        """
        class M(nn.Module):
            def forward(self, x):
                x = x + x
                return x

        m = M()
        qconfig_dict = {
            '': None,
            'object_type': [
                (torch.add, torch.quantization.default_qconfig),
            ],
        }
        example_args = (torch.randn(1, 1, 1, 1),)
        mp = _quantize_dbr.prepare(m, qconfig_dict, example_args)
        mp(*example_args)
        mq = _quantize_dbr.convert(mp)
        mq(*example_args)
        rewritten = mq.rewrite_for_scripting()
        expected_occurrence = {
            NodeSpec.call_function(torch.add): 0,
            NodeSpec.call_function(toq.add): 1,
        }
        self.checkGraphModuleNodes(
            rewritten, expected_node_occurrence=expected_occurrence)

    def test_qconfig_dict_module_name(self):
        """
        Verifies that the 'module_name' option of qconfig_dict works
        on module types.
        """
        m = nn.Sequential(
            nn.Sequential(
                nn.Conv2d(1, 1, 1),
            ),
            nn.Conv2d(1, 1, 1),
            nn.Sequential(
                nn.Conv2d(1, 1, 1),
                nn.Conv2d(1, 1, 1),
            ),
        )
        qconfig_dict = {
            '': torch.quantization.default_qconfig,
            'module_name': [
                ('0', torch.quantization.default_qconfig),
                ('1', None),
                ('2.0', None),
            ],
        }
        example_args = (torch.randn(1, 1, 1, 1),)
        mp = _quantize_dbr.prepare(m, qconfig_dict, example_args)
        mp(*example_args)
        mq = _quantize_dbr.convert(mp)
        mq(*example_args)
        self.assertTrue(isinstance(mq[0][0], nnq.Conv2d))
        self.assertTrue(isinstance(mq[1], nn.Conv2d))
        self.assertTrue(isinstance(mq[2][0], nn.Conv2d))
        self.assertTrue(isinstance(mq[2][1], nnq.Conv2d))

    def test_qconfig_dict_unsupported_does_not_crash_when_empty(self):
        """
        Verifies that the yet unimplemented keys of qconfig_dict only
        crash when they have non-zero values
        """
        m = nn.Sequential(nn.Conv2d(1, 1, 1)).eval()
        qconfig_dict = {'': torch.quantization.default_qconfig}
        # this modifies qconfig_dict inplace to include more keys
        mp = prepare_fx(m, qconfig_dict)
        example_args = (torch.randn(1, 1, 1, 1),)
        # need this line to not crash
        mp = _quantize_dbr.prepare(m, qconfig_dict, example_args)

    def _test_serialization(self, model, input_shape):
        example_inputs = (torch.randn(*input_shape),)
        qconfig_dict = {'': torch.quantization.default_qconfig}
        m = model().eval()
        m = _quantize_dbr.prepare(m, qconfig_dict, example_inputs)
        # calibrate, to populate statistics
        m(example_inputs[0])
        m = _quantize_dbr.convert(m)

        qconfig_dict = {'': torch.quantization.default_qconfig}
        m2 = model().eval()
        m2 = _quantize_dbr.prepare(m2, qconfig_dict, example_inputs)
        # do not calibrate, to ensure important statistics are populated without calibration and
        # the results are different at every node, including the quantize_per_tensor node
        m2 = _quantize_dbr.convert(m2)

        # Results should be different without loading from serialized state_dict
        expected = m(example_inputs[0])
        actual = m2(example_inputs[0])
        self.assertFalse(_allclose(expected, actual))

        # Results should be the same after loading from serialized state_dict
        with tempfile.NamedTemporaryFile(delete=False) as f:
            torch.save(m.state_dict(), f)
            with open(f.name, 'rb') as f2:
                loaded_state_dict = torch.load(f2)
                m2.load_state_dict(loaded_state_dict)
        expected = m(example_inputs[0])
        actual = m2(example_inputs[0])
        self.assertTrue(_allclose(expected, actual))

    def test_serialization(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)
                self.linear = torch.nn.Linear(1, 1)

            def forward(self, x):
                x1 = self.conv(x)
                x2 = self.linear(x1)
                return x2

        input_shape = (1, 1, 1, 1)
        self._test_serialization(M, input_shape)

    def test_serialization_functional(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                # conv
                self.weight2d = torch.nn.Parameter(torch.randn(1, 1, 1, 1))
                self.bias2d = torch.nn.Parameter(torch.randn(1))
                self.stride2d = (1, 1)
                self.padding2d = (0, 0)
                self.dilation2d = (1, 1)
                self.groups = 1
                # linear
                self.w1 = nn.Parameter(torch.empty(1, 1))
                self.b1 = nn.Parameter(torch.ones(1))
                torch.nn.init.kaiming_uniform_(self.w1, a=math.sqrt(5))

            def forward(self, x):
                updated_weight = self.weight2d * x
                x = F.conv2d(
                    x, updated_weight, self.bias2d, self.stride2d, self.padding2d,
                    self.dilation2d, self.groups)
                # TODO: Investigate why serialization does not work with functional linear
                # x = F.linear(x, self.w1, self.b1)
                return x

        input_shape = (1, 1, 1, 1)
        self._test_serialization(M, input_shape)

    def test_jit_tracing_removes_aliases(self):
        m = nn.Sequential(
            nn.Conv2d(1, 1, 1),
            nn.Sequential(
                nn.Conv2d(1, 1, 1),
            ),
        )
        qconfig_dict = {'': torch.quantization.default_qconfig}
        example_args = (torch.randn(1, 1, 1, 1),)
        mp = _quantize_dbr.prepare(m, qconfig_dict, example_args)
        mq = _quantize_dbr.convert(mp)
        mqs = torch.jit.trace(mq, example_args)
        FileCheck().check_count("aten::alias", 5, exactly=True).run(
            mqs.inlined_graph)
        res1 = mqs(*example_args)
        mqs = remove_redundant_aliases(mqs)
        res2 = mqs(*example_args)
        self.assertTrue(torch.allclose(res1, res2))
        # TODO(future PR): figure out why aliasing still appears in the inlined
        # graph, and if that is fixed then just check the inlined graph.
        for graph in (
            mqs.graph,
            getattr(mqs, '1').graph,
            getattr(getattr(mqs, '1'), '0').graph,
        ):
            FileCheck().check_count("aten::alias", 0, exactly=True).run(graph)

    def test_conv_int32_reference_model(self):
        m = nn.Sequential(nn.Conv2d(1, 1, 1)).eval()
        int32_obs_ctr = MinMaxObserver.with_args(dtype=torch.qint32)
        int32_qconfig = QConfig(weight=int32_obs_ctr, activation=int32_obs_ctr)
        qconfig_dict = {'': int32_qconfig}
        mp = _quantize_dbr.prepare(m, qconfig_dict, (torch.randn(1, 1, 1, 1),))
        mp(torch.randn(1, 1, 1, 1))
        mq = _quantize_dbr.convert(mp)
        res = mq(torch.randn(1, 1, 1, 1))
        mqt = torch.jit.trace(mq, (torch.randn(1, 1, 1, 1),))
        # verify the right ops are present:
        # x0 -> quant -> (dequant -> conv_ref -> quant) -> dequant -> x1
        FileCheck()\
            .check_count("aten::quantize_per_tensor", 2, exactly=True)\
            .run(mqt.graph)
        FileCheck()\
            .check_count("aten::dequantize", 2, exactly=True)\
            .run(mqt.graph)

@skipIfNoFBGEMM
class TestQuantizeDBRMultipleOps(QuantizeDBRTestCase):
    """
    Tests that DBR quantization covers interactions between multiple ops

    Most of these tests were added when the code was an early prototype
    and were one-off test cases for patterns which happened to break
    the code on various models at the time of writing. A lot of these
    can probably be removed in the future as they are replaced by more
    systematic individual and fusion tests.
    """
    def test_dropout_conv(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.dropout = nn.Dropout()
                self.conv = torch.nn.Conv2d(1, 1, 1)

            def forward(self, x):
                # this can be sometimes inplace
                x1 = self.dropout(x)
                x1 = self.conv(x)
                return x1

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_conv_flatten_linear(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)
                self.linear = torch.nn.Linear(1, 1)

            def forward(self, x):
                x1 = self.conv(x)
                # TODO(future PR): unbreak this
                # x1 = torch.nn.functional.adaptive_avg_pool2d(x, (1, 1))
                x1 = torch.nn.functional.adaptive_avg_pool2d(x1, (1, 1))
                x2 = torch.flatten(x1, 1)
                x3 = self.linear(x2)
                return x3

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 1, 1),))

    def test_conv_add(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)

            def forward(self, x):
                x1 = self.conv(x)
                x2 = x1 + x
                return x2

        m = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(m, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_conv_scalar_add(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)

            def forward(self, x):
                x = self.conv(x)
                x = x + 1.0
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_conv_relu_add(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = torch.nn.Conv2d(1, 1, 1)
                self.relu = torch.nn.ReLU()

            def forward(self, x):
                x1 = self.conv(x)
                x2 = self.relu(x1)
                x3 = x1 + x
                return x3

        model_fp32 = M().eval()

        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 2, 2),))

    def test_linear_torch_relu(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.u1 = nn.Linear(1, 1)
                self.v1 = nn.Linear(1, 1)
                self.u2 = nn.Linear(1, 1)
                self.v2 = nn.Linear(1, 1)
                self.w = nn.Linear(1, 1)

            def forward(self, x):
                x = self.w(x)
                x = x + torch.relu(self.v1(torch.relu(self.u1(x))))
                return x + torch.relu(self.v2(torch.relu(self.u2(x))))

        model_fp32 = M().eval()

        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 1, 1),))

    def test_gelu_linear(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.gelu = torch.nn.GELU()
                self.linear = torch.nn.Linear(1, 1)

            def forward(self, x):
                x = self.linear(x)
                x = self.gelu(x)
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 1, 1),))

    def test_dropout(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.dropout = nn.Dropout()
                self.linear = torch.nn.Linear(1, 1)
                self.linear2 = torch.nn.Linear(1, 1)

            def forward(self, x):
                x = self.linear(x)
                x = self.dropout(x)
                x = self.linear2(x)
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 1, 1),))

    def test_module_then_add(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(1, 1)

            def forward(self, x):
                x = self.linear(x)
                x = x + 1.0
                x = x + 1.0
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 1, 1),))

    def test_add_linear(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = nn.Linear(1, 1)

            def forward(self, x):
                x = x + x
                x = self.linear(x)
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(model_fp32, qconfig, (torch.randn(1, 1, 1, 1),))

    def test_inplace_add(self):
        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.embedding1 = nn.Embedding(1, 1)
                self.embedding2 = nn.Embedding(1, 1)
                self.layernorm = nn.LayerNorm(1)

            def forward(self, x):
                x1 = self.embedding1(x)
                x1 += self.embedding2(x)
                x2 = self.layernorm(x1)
                return x

        model_fp32 = M().eval()
        qconfig = torch.quantization.default_qconfig
        prepare_custom_config_dict = {
            'output_dtypes': (torch.int64,),
        }
        self._test_auto_tracing(
            model_fp32, qconfig, (torch.LongTensor([[0]]),),
            fuse_modules=False,
            dbr_prepare_custom_config_dict=prepare_custom_config_dict)

    def test_lstm(self):
        # building block of torchbenchmark/tts_angular
        class LSTMWithProjection(nn.Module):
            def __init__(self, input_size, hidden_size, proj_size):
                super().__init__()
                self.input_size = input_size
                self.hidden_size = hidden_size
                self.proj_size = proj_size
                self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
                self.linear = nn.Linear(hidden_size, proj_size, bias=False)

            def forward(self, x):
                self.lstm.flatten_parameters()
                o, (_, _) = self.lstm(x)
                return self.linear(o)

        m = LSTMWithProjection(1, 1, 1).eval()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            m, qconfig, (torch.randn(1, 1, 1),),
            # the module is not symbolically traceable
            do_fx_comparison=False)


@skipIfNoFBGEMM
class TestQuantizeDBRModels(QuantizeDBRTestCase):
    @skip_if_no_torchvision
    def test_mobilenet_v2(self):
        import torchvision
        m = torchvision.models.__dict__['mobilenet_v2'](pretrained=False).eval().float()
        qconfig = torch.quantization.default_qconfig
        self._test_auto_tracing(
            m, qconfig, (torch.randn(1, 3, 224, 224),),
            # TODO fix this (reason TBD)
            do_torchscript_checks=False)

    @skip_if_no_torchvision
    def test_mobilenet_v2_removes_aliases(self):
        import torchvision
        m = torchvision.models.__dict__['mobilenet_v2'](pretrained=False)\
            .eval().float()
        qconfig_dict = {'': torch.quantization.default_qconfig}
        example_args = (torch.randn(1, 3, 224, 224),)
        mp = _quantize_dbr.prepare(m, qconfig_dict, example_args)
        mq = _quantize_dbr.convert(mp)
        mqs = torch.jit.trace(mq, example_args)
        res1 = mqs(*example_args)
        mqs = remove_redundant_aliases(mqs)
        res2 = mqs(*example_args)
        self.assertTrue(torch.allclose(res1, res2))
