Repository: quic/aimet (language: python)

Issue:
Support reduce ops and elementwise boolean



Discussion / hints:

Fix patch:
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/elementwise_ops.py b/TrainingExtensions/torch/src/python/aimet_torch/elementwise_ops.py
index df251e6b70..e3284c3cef 100644
--- a/TrainingExtensions/torch/src/python/aimet_torch/elementwise_ops.py
+++ b/TrainingExtensions/torch/src/python/aimet_torch/elementwise_ops.py
@@ -92,6 +92,9 @@ def create_wrapper_module(class_name: str, functional: Callable) -> Callable:
 Bmm = create_wrapper_module('Bmm', torch.bmm)
 CumSum = create_wrapper_module('CumSum', torch.cumsum)
 MaskedFill = create_wrapper_module('MaskedFill', torch.Tensor.masked_fill_)
+Mean = create_wrapper_module('Mean', torch.mean)
+Sum = create_wrapper_module('Sum', torch.sum)
+Prod = create_wrapper_module('Prod', torch.prod)
 
 
 # modules for functional operations defined under torch.nn.functional package
@@ -101,7 +104,6 @@ def create_wrapper_module(class_name: str, functional: Callable) -> Callable:
 AvgPool2d = create_wrapper_module('AvgPool2d', torch.nn.functional.avg_pool2d)
 BatchNorm = create_wrapper_module('BatchNorm', torch.nn.functional.batch_norm)
 GroupNorm = create_wrapper_module('GroupNorm', torch.nn.functional.group_norm)
-Mean = create_wrapper_module('Mean', torch.mean)
 
 
 # following modules are for overloaded operators like + and *,


Test patch:
diff --git a/TrainingExtensions/torch/test/python/test_elementwise_ops.py b/TrainingExtensions/torch/test/python/test_elementwise_ops.py
index d9a0bb8664..f3bb1cb8c7 100644
--- a/TrainingExtensions/torch/test/python/test_elementwise_ops.py
+++ b/TrainingExtensions/torch/test/python/test_elementwise_ops.py
@@ -44,7 +44,7 @@
 import aimet_common.libpymo as libpymo
 from aimet_common.defs import QuantScheme
 from aimet_torch.elementwise_ops import Add, Subtract, Multiply, Divide, Concat, MatMul, Erf, Sqrt, \
-    Greater, Less, GreaterEqual, LessEqual, NotEqual, Equal, Where
+    Greater, Less, GreaterEqual, LessEqual, NotEqual, Equal, Where, Mean, Sum, Prod
 from aimet_torch.quantsim import QuantizationSimModel
 
 
@@ -413,3 +413,27 @@ def test_where_op(self):
         out1 = torch.where(input1 > input2, input1, input2)
         self.assertTrue(np.allclose(out, out1))
         self.assertEqual(input1.shape, out.shape)
+
+    def test_mean_op(self):
+        torch.manual_seed(10)
+        model = Model3(Mean())
+        input1 = torch.rand((5, 10, 10, 20))
+        out = model(input1)
+        out1 = torch.mean(input1)
+        self.assertTrue(np.allclose(out, out1))
+
+    def test_sum_op(self):
+        torch.manual_seed(10)
+        model = Model3(Sum())
+        input1 = torch.rand((5, 10, 10, 20))
+        out = model(input1)
+        out1 = torch.sum(input1)
+        self.assertTrue(np.allclose(out, out1))
+
+    def test_prod_op(self):
+        torch.manual_seed(10)
+        model = Model3(Prod())
+        input1 = torch.rand((5, 10, 10, 20))
+        out = model(input1)
+        out1 = torch.prod(input1)
+        self.assertTrue(np.allclose(out, out1))