pytorch · karljang · Apr 2, 2019 · Apr 3, 2019 · Apr 3, 2019 · Apr 8, 2019
diff --git a/test/onnx/expect/TestOperators.test_maxpool_dilations.expect b/test/onnx/expect/TestOperators.test_maxpool_dilations.expect
@@ -0,0 +1,73 @@
+ir_version: 4
+producer_name: "pytorch"
+producer_version: "1.1"
+graph {
+  node {
+    input: "0"
+    output: "1"
+    op_type: "MaxPool"
+    attribute {
+      name: "dilations"
+      ints: 2
+      type: INTS
+    }
+    attribute {
+      name: "kernel_shape"
+      ints: 2
+      type: INTS
+    }
+    attribute {
+      name: "pads"
+      ints: 0
+      ints: 0
+      type: INTS
+    }
+    attribute {
+      name: "strides"
+      ints: 1
+      type: INTS
+    }
+  }
+  name: "torch-jit-export"
+  input {
+    name: "0"
+    type {
+      tensor_type {
+        elem_type: 1
+        shape {
+          dim {
+            dim_value: 20
+          }
+          dim {
+            dim_value: 16
+          }
+          dim {
+            dim_value: 50
+          }
+        }
+      }
+    }
+  }
+  output {
+    name: "1"
+    type {
+      tensor_type {
+        elem_type: 1
+        shape {
+          dim {
+            dim_value: 20
+          }
+          dim {
+            dim_value: 16
+          }
+          dim {
+            dim_value: 48
+          }
+        }
+      }
+    }
+  }
+}
+opset_import {
+  version: 10
+}
diff --git a/test/onnx/test_operators.py b/test/onnx/test_operators.py
@@ -258,6 +258,10 @@ def test_maxpool(self):
         x = torch.randn(20, 16, 50)
         self.assertONNX(nn.MaxPool1d(3, stride=2), x)
 
+    def test_maxpool_dilations(self):
+        x = torch.randn(20, 16, 50)
+        self.assertONNX(nn.MaxPool1d(2, stride=1, dilation=2), x, opset_version=10)
+
     def test_avg_pool2d(self):
         x = torch.randn(20, 16, 50, 32)
         self.assertONNX(nn.AvgPool2d(3, stride=2), x)

diff --git a/torch/onnx/symbolic.py b/torch/onnx/symbolic.py
@@ -672,8 +672,6 @@ def _max_pool(name, tuple_fn, ndims, return_indices):
     def symbolic_fn(g, input, kernel_size, stride, padding, dilation, ceil_mode):
         if ceil_mode and input.type().kind() != "CompleteTensorType":
             return _unimplemented(name, "input size not accesible")
-        if set(tuple_fn(dilation)) != {1}:
-            return _unimplemented(name, "dilation")
         if not stride:
             stride = kernel_size
         padding = tuple(tuple_fn(padding))
@@ -682,6 +680,13 @@ def symbolic_fn(g, input, kernel_size, stride, padding, dilation, ceil_mode):
             padding = padding + tuple(numpy.add(padding_ceil, padding))
         else:
             padding = padding * 2
+        kwargs = {
+            'kernel_shape_i': tuple_fn(kernel_size),
+            'pads_i': padding,
+            'strides_i': tuple_fn(stride),
+        }
+        if set(tuple_fn(dilation)) != {1}:
+            kwargs['dilations_i'] = tuple_fn(dilation)
         # easy but hacky way to get flattened indices values
         # to be used to convert the indices values to non-flattened.
         # In ONNX the indices are computed as a flatten 1-D tensor,
@@ -696,10 +701,7 @@ def symbolic_fn(g, input, kernel_size, stride, padding, dilation, ceil_mode):
         # For more information :
         # https://github.com/pytorch/pytorch/pull/16455#issuecomment-460776407
         if return_indices:
-            r, indices = g.op("MaxPool", input, outputs=2,
-                              kernel_shape_i=tuple_fn(kernel_size),
-                              pads_i=padding,
-                              strides_i=tuple_fn(stride))
+            r, indices = g.op("MaxPool", input, outputs=2, **kwargs)
             _, flattened_indices = g.op("MaxPool", input, outputs=2,
                                         kernel_shape_i=[1 for _ in range(ndims)],
                                         strides_i=[1 for _ in range(ndims)])
@@ -709,10 +711,7 @@ def symbolic_fn(g, input, kernel_size, stride, padding, dilation, ceil_mode):
             indices = sub(g, indices, s)
             return r, indices
         else:
-            r = g.op("MaxPool", input, outputs=1,
-                     kernel_shape_i=tuple_fn(kernel_size),
-                     pads_i=padding,
-                     strides_i=tuple_fn(stride))
+            r = g.op("MaxPool", input, outputs=1, **kwargs)
             return r
 
     return symbolic_fn