From 2a46c152f67e755e99212b1e886a183954f7642a Mon Sep 17 00:00:00 2001
From: Christoph Gerum <christoph.gerum@uni-tuebingen.de>
Date: Thu, 18 Mar 2021 11:59:41 +0100
Subject: [PATCH] [TORCH] Implement avg_pool1d

---
 python/tvm/relay/frontend/pytorch.py          | 25 +++++++++++++++++++
 tests/python/frontend/pytorch/test_forward.py | 19 ++++++++++++--
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
index fd0a07e35c155..82df97cefd053 100644
--- a/python/tvm/relay/frontend/pytorch.py
+++ b/python/tvm/relay/frontend/pytorch.py
@@ -1353,6 +1353,30 @@ def softplus(self, inputs, input_types):
         beta = _expr.const(float(inputs[1]), dtype=dtype)
         return _op.log(_op.exp(inputs[0] * beta) + _expr.const(1.0, dtype=dtype)) / beta
 
+    def avg_pool1d(self, inputs, input_types):
+        data = inputs[0]
+
+        pool_size = self.convert_const_list(inputs[1])
+        strides = self.convert_const_list(inputs[2] if inputs[2] else pool_size)
+        padding = inputs[3]
+        ceil_mode = int(inputs[4])
+        count_include_pad = int(inputs[5])
+
+        def func(x):
+            return _op.nn.avg_pool1d(
+                x,
+                pool_size=pool_size,
+                strides=strides,
+                padding=padding,
+                ceil_mode=ceil_mode,
+                count_include_pad=count_include_pad,
+            )
+
+        if self.is_quantized_tensor(data):
+            return qnn_torch.apply_with_upcast(data, func)
+
+        return func(data)
+
     def avg_pool2d(self, inputs, input_types):
         data = inputs[0]
 
@@ -2338,6 +2362,7 @@ def create_convert_map(self):
             "aten::log_softmax": self.log_softmax,
             "aten::sigmoid": self.sigmoid,
             "aten::softplus": self.softplus,
+            "aten::avg_pool1d": self.avg_pool1d,
             "aten::avg_pool2d": self.avg_pool2d,
             "aten::avg_pool3d": self.avg_pool3d,
             "aten::linear": self.linear,
diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py
index 83c1698799c73..e02b89d27162f 100644
--- a/tests/python/frontend/pytorch/test_forward.py
+++ b/tests/python/frontend/pytorch/test_forward.py
@@ -809,7 +809,21 @@ def forward(self, *args):
 
 
 @tvm.testing.uses_gpu
-def test_forward_avgpool():
+def test_forward_avgpool1d():
+    torch.set_grad_enabled(False)
+    input_shape = [1, 3, 10]
+
+    class AvgPool1D2(Module):
+        def forward(self, *args):
+            return torch.nn.functional.avg_pool1d(args[0], kernel_size=[10])
+
+    input_data = torch.rand(input_shape).float()
+    verify_model(torch.nn.AvgPool1d(kernel_size=[10]).eval(), input_data=input_data)
+    verify_model(AvgPool1D2().float().eval(), input_data=input_data)
+
+
+@tvm.testing.uses_gpu
+def test_forward_avgpool2d():
     torch.set_grad_enabled(False)
     input_shape = [1, 3, 10, 10]
 
@@ -3831,7 +3845,8 @@ def test_fn(is_sorted, return_inverse, return_counts):
     test_forward_logsoftmax()
     test_forward_sigmoid()
     test_forward_dense()
-    test_forward_avgpool()
+    test_forward_avgpool1d()
+    test_forward_avgpool2d()
     test_forward_avgpool3d()
     test_forward_dropout()
     test_forward_slice()