pytorch · goldsborough · Jun 26, 2018 · Jun 25, 2018 · Jun 26, 2018
diff --git a/aten/src/ATen/native/Pooling.cpp b/aten/src/ATen/native/Pooling.cpp
@@ -1,25 +1,26 @@
 #include "ATen/ATen.h"
-#include "ATen/TensorUtils.h"
-#include "ATen/NativeFunctions.h"
 
-#include <sstream>
-#include <vector>
+#include "ATen/Error.h"
+#include "ATen/NativeFunctions.h"
+#include "ATen/TensorUtils.h"
 
+#include <tuple>
 
 namespace at { namespace native {
 
-static void check1d(const char* name, IntList x) {
-  if (x.size() != 1) {
-    std::ostringstream ss;
-    ss << "max_pool1d() argument '" << name << "' should contain one int (got "
-       << x.size() << ")";
-    throw std::runtime_error(ss.str());
-  }
+static void check1d(
+    const char* function_name,
+    const char* argument_name,
+    IntList x) {
+  AT_CHECK(
+      x.size() == 1,
+      function_name, "() argument '", argument_name,
+      "' should contain one int (got ", x.size(), ")");
 }
 
 Tensor adaptive_avg_pool1d(const Tensor & self, IntList output_size) {
   checkDim("adaptive_avg_pool1d", TensorArg(self, "self", 1), 3);
-  check1d("output_size", output_size);
+  check1d("adaptive_avg_pool1d", "output_size", output_size);
 
   auto output = at::adaptive_avg_pool2d(
       self.unsqueeze(2),
@@ -30,7 +31,7 @@ Tensor adaptive_avg_pool1d(const Tensor & self, IntList output_size) {
 
 std::tuple<Tensor,Tensor> adaptive_max_pool1d(const Tensor & self, IntList output_size) {
   checkDim("adaptive_max_pool1d", TensorArg(self, "self", 1), 3);
-  check1d("output_size", output_size);
+  check1d("adaptive_max_pool1d", "output_size", output_size);
 
   Tensor output, indices;
   std::tie(output, indices) = at::adaptive_max_pool2d(
@@ -48,10 +49,10 @@ std::tuple<Tensor,Tensor> max_pool1d(
     stride = kernel_size;
   }
   checkDim("max_pool1d", TensorArg(self, "self", 1), 3);
-  check1d("kernel_size", kernel_size);
-  check1d("stride", stride);
-  check1d("padding", padding);
-  check1d("dilation", dilation);
+  check1d("max_pool1d", "kernel_size", kernel_size);
+  check1d("max_pool1d", "stride", stride);
+  check1d("max_pool1d", "padding", padding);
+  check1d("max_pool1d", "dilation", dilation);
 
   Tensor output, indices;
   std::tie(output, indices) = at::max_pool2d(
@@ -65,4 +66,30 @@ std::tuple<Tensor,Tensor> max_pool1d(
   return std::make_tuple(output.squeeze(2), indices.squeeze(2));
 }
 
-}}  // namespace at::native
+Tensor avg_pool1d(
+    const Tensor& self,
+    IntList kernel_size,
+    IntList stride,
+    IntList padding,
+    bool ceil_mode,
+    bool count_include_pad) {
+  if (stride.empty()) {
+    stride = kernel_size;
+  }
+  checkDim("avg_pool1d", TensorArg(self, "self", 1), 3);
+  check1d("avg_pool1d", "kernel_size", kernel_size);
+  check1d("avg_pool1d", "stride", stride);
+  check1d("avg_pool1d", "padding", padding);
+
+  auto output = at::avg_pool2d(
+      self.unsqueeze(2),
+      {1, kernel_size[0]},
+      {1, stride[0]},
+      {0, padding[0]},
+      ceil_mode,
+      count_include_pad);
+
+  return output.squeeze(2);
+}
+} // namespace native
+} // namespace at
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -75,6 +75,9 @@
     CPU: _acos_out_cpu
     CUDA: _acos_out_cuda
 
+- func: avg_pool1d(Tensor self, IntList[1] kernel_size, IntList[1] stride={}, IntList[1] padding=0, bool ceil_mode=false, bool count_include_pad=true) -> Tensor
+  variants: function
+
 - func: adaptive_avg_pool1d(Tensor self, IntList[1] output_size) -> Tensor
   variants: function
 

diff --git a/torch/nn/functional.py b/torch/nn/functional.py
@@ -212,40 +212,33 @@ def conv_tbc(input, weight, bias, pad=0):
 
 
 # Pooling
-def avg_pool1d(input, kernel_size, stride=None, padding=0,
-               ceil_mode=False, count_include_pad=True):
-    r"""Applies a 1D average pooling over an input signal composed of several
-    input planes.
+avg_pool1d = _add_docstr(torch.avg_pool1d, r"""
+avg_pool1d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True) -> Tensor
 
-    See :class:`~torch.nn.AvgPool1d` for details and output shape.
+Applies a 1D average pooling over an input signal composed of several
+input planes.
 
-    Args:
-        input: input tensor of shape (:math:`minibatch \times in\_channels \times iW`)
-        kernel_size: the size of the window. Can be a single number or a
-          tuple `(kW,)`
-        stride: the stride of the window. Can be a single number or a tuple
-          `(sW,)`. Default: :attr:`kernel_size`
-        padding: implicit zero paddings on both sides of the input. Can be a
-          single number or a tuple `(padW,)`. Default: 0
-        ceil_mode: when True, will use `ceil` instead of `floor` to compute the
-            output shape. Default: ``False``
-        count_include_pad: when True, will include the zero-padding in the
-            averaging calculation. Default: ``True``
+See :class:`~torch.nn.AvgPool1d` for details and output shape.
 
-    Example::
-        >>> # pool of square window of size=3, stride=2
-        >>> input = torch.tensor([[[1,2,3,4,5,6,7]]])
-        >>> F.avg_pool1d(input, kernel_size=3, stride=2)
-        tensor([[[ 2.,  4.,  6.]]])
-    """
-    if input.dim() != 3:
-        raise ValueError('expected 3D input (got {} dimensions)'
-                         .format(input.dim()))
-    kernel_size = _single(kernel_size) + (1,)
-    stride = _single(stride) + (1,) if stride is not None else kernel_size
-    padding = _single(padding) + (0,)
-    return avg_pool2d(input.unsqueeze(3), kernel_size, stride, padding,
-                      ceil_mode, count_include_pad).squeeze(3)
+Args:
+    input: input tensor of shape (:math:`minibatch \times in\_channels \times iW`)
+    kernel_size: the size of the window. Can be a single number or a
+      tuple `(kW,)`
+    stride: the stride of the window. Can be a single number or a tuple
+      `(sW,)`. Default: :attr:`kernel_size`
+    padding: implicit zero paddings on both sides of the input. Can be a
+      single number or a tuple `(padW,)`. Default: 0
+    ceil_mode: when True, will use `ceil` instead of `floor` to compute the
+        output shape. Default: ``False``
+    count_include_pad: when True, will include the zero-padding in the
+        averaging calculation. Default: ``True``
+
+Example::
+    >>> # pool of square window of size=3, stride=2
+    >>> input = torch.tensor([[[1,2,3,4,5,6,7]]])
+    >>> F.avg_pool1d(input, kernel_size=3, stride=2)
+    tensor([[[ 2.,  4.,  6.]]])
+""")
 
 
 avg_pool2d = _add_docstr(torch._C._nn.avg_pool2d, r"""