Test if quantized modules are torch.jit traceable

TrainingExtensions/torch/src/python/aimet_torch/v2/nn/true_quant.py

@@ -874,7 +874,9 @@ class QuantizedFeatureAlphaDropout(_DispatchMixin, QuantizationMixin, nn.Feature
 @QuantizationMixin.implements(nn.Flatten)
 class QuantizedFlatten(_DispatchMixin, QuantizationMixin, nn.Flatten):
     """ Quantized Flatten """
-    _builtin_torch_fn = Tensor.flatten
+    def _get_builtin_torch_fn(self):
+        return Tensor.flatten
+
     __quant_init__ = __unary__
 
 
@@ -1836,7 +1838,8 @@ class QuantizedTripletMarginWithDistanceLoss(_DispatchMixin, QuantizationMixin,
 @QuantizationMixin.implements(nn.Unflatten)
 class QuantizedUnflatten(_DispatchMixin, QuantizationMixin, nn.Unflatten):
     """ Quantized Unflatten """
-    _builtin_torch_fn = Tensor.unflatten
+    def _get_builtin_torch_fn(self):
+        return Tensor.unflatten
 
 
 @QuantizationMixin.implements(nn.Unfold)

TrainingExtensions/torch/test/python/v2/nn/test_true_quant.py

@@ -41,10 +41,10 @@
 
 import pytest
 import torch
-from torch import randn, randint, zeros, full, arange, ones
+from torch import randn, randint, zeros, full, arange, ones, tensor
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.utils._pytree import tree_map, tree_flatten
+from torch.utils._pytree import tree_flatten
 from torch.overrides import get_ignored_functions
 from aimet_torch.v2.quantization.affine.backends import quantize, quantize_dequantize, dequantize
 from aimet_torch.v2.quantization.affine import Quantize, QuantizeDequantize
@@ -67,7 +67,7 @@
 from aimet_torch.v2.nn.fake_quant import _legacy_impl
 from aimet_torch.v2.nn.true_quant import _dispatch, _dispatch_table
 from aimet_torch.v2.quantization.affine import AffineEncoding
-from aimet_torch.v2.quantization.tensor import QuantizedTensorBase, QuantizedTensor, DequantizedTensor
+from aimet_torch.v2.quantization.tensor import QuantizedTensor, DequantizedTensor
 from aimet_torch.v2.utils import enable_recompute
 from aimet_torch.v2.nn import custom
 
@@ -886,16 +886,16 @@ def _create_quantized_module(module):
 [
     (lambda: custom.Sin(),                            lambda: randn(100)),
     (lambda: custom.Cos(),                            lambda: randn(100)),
-    (lambda: custom.AvgPool2d(),                      lambda: (randn(1,10,10), 2)),
-    (lambda: custom.Reshape(),                        lambda: (randn(10,10), (100, 1))),
+    (lambda: custom.AvgPool2d(),                      lambda: (randn(1,10,10), (tensor(2),))),
+    (lambda: custom.Reshape(),                        lambda: (randn(10,10), (tensor(100), tensor(1)))),
     (lambda: custom.RSqrt(),                          lambda: randn(100).abs()),
     (lambda: custom.Add(),                            lambda: (randn(100), randn(100))),
     (lambda: custom.Multiply(),                       lambda: (randn(100), randn(100))),
     (lambda: custom.Subtract(),                       lambda: (randn(100), randn(100))),
     (lambda: custom.Divide(),                         lambda: (randn(100), randn(100))),
     (lambda: custom.Concat(),                         lambda: (randn(1, 100), randn(3, 100))),
 ]))
-def test_default_kernel_abtest(module_factory, input_factory):
+def test_default_kernels(module_factory, input_factory):
     module = module_factory()
     inputs = input_factory()
 
@@ -926,6 +926,18 @@ def test_default_kernel_abtest(module_factory, input_factory):
     torch.manual_seed(0)
     out = qmodule(*inputs)
 
-    for out, fout in zip(tree_flatten(out)[0], tree_flatten(fout)[0]):
-        assert torch.equal(out, fout)
-        assert torch.all(out.isfinite())
+    for out_, fout_ in zip(tree_flatten(out)[0], tree_flatten(fout)[0]):
+        assert torch.equal(out_, fout_)
+        assert torch.all(out_.isfinite())
+
+    """
+    When: Trace a quantized modules with torch.jit.trace
+    Then: 1) Tracing shouldn't fail
+          2) The traced module should produce the same output as the original module
+    """
+    traced = torch.jit.trace(qmodule, inputs)
+    torch.manual_seed(0)
+    tout = traced(*inputs)
+
+    for out_, tout_ in zip(tree_flatten(out)[0], tree_flatten(tout)[0]):
+        assert torch.equal(out_, tout_)