From 6726971625806781a56ba4ba9221a1519ff01f2f Mon Sep 17 00:00:00 2001
From: Huyuwei <huyuwei1995@gmail.com>
Date: Sun, 13 Aug 2017 23:23:28 +0800
Subject: [PATCH 1/5] add dilation operators

---
 topi/python/topi/nn/__init__.py |  1 +
 topi/python/topi/nn/dilate.py   | 63 +++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)
 create mode 100644 topi/python/topi/nn/dilate.py

diff --git a/topi/python/topi/nn/__init__.py b/topi/python/topi/nn/__init__.py
index 80366e09f402..884989920c60 100644
--- a/topi/python/topi/nn/__init__.py
+++ b/topi/python/topi/nn/__init__.py
@@ -5,3 +5,4 @@
 from .mapping import *
 from .ewise import *
 from .conv import *
+from .dilate import *
diff --git a/topi/python/topi/nn/dilate.py b/topi/python/topi/nn/dilate.py
new file mode 100644
index 000000000000..c06905c69dba
--- /dev/null
+++ b/topi/python/topi/nn/dilate.py
@@ -0,0 +1,63 @@
+"""Dilation operators"""
+from __future__ import absolute_import as _abs
+from .util import get_const_tuple
+import tvm
+
+@tvm.tag_scope(tag="dilation")
+def dilate_nchw(Input, stride_h, stride_w):
+    """Dilate Input with zeros.
+
+    Parameters
+    ----------
+    Input : tvm.Tensor
+        Input tensor, layout is NCHW.
+
+    stride_h : int
+        Dilation extent on H dimension, 1 means no dilation.
+
+    stride_w : int
+        Dilation extent on W dimension, 1 means no dilation.
+
+    Returns
+    -------
+    Output : tvm.Tensor
+        Output tensor, layout is NCHW.
+    """
+    N, C, H, W = get_const_tuple(Input.shape)
+    Output = tvm.compute(
+        (N, C, (H-1)*stride_h+1, (W-1)*stride_w+1),
+        lambda n, c, h, w: tvm.select(
+            tvm.all(tvm.make.EQ(h%stride_h, 0), tvm.make.EQ(w%stride_w, 0)),
+            Input(n, c, h/stride_h, w/stride_w), tvm.const(0.0)),
+        name='DilatedInput')
+    return Output
+
+
+@tvm.tag_scope(tag="dilation")
+def dilate_nhwc(Input, stride_h, stride_w):
+    """Dilate Input with zeros.
+
+    Parameters
+    ----------
+    Input : tvm.Tensor
+        Input tensor, layout is NHWC.
+
+    stride_h : int
+        Dilation extent on H dimension, 1 means no dilation.
+
+    stride_w : int
+        Dilation extent on W dimension, 1 means no dilation.
+
+    Returns
+    -------
+    Output : tvm.Tensor
+        Output tensor, layout is NHWC.
+    """
+    N, H, W, C = get_const_tuple(Input.shape)
+    Output = tvm.compute(
+        (N, (H-1)*stride_h+1, (W-1)*stride_w+1, C),
+        lambda n, h, w, c: tvm.select(
+            tvm.all(tvm.make.EQ(h%stride_h, 0), tvm.make.EQ(w%stride_w, 0)),
+            Input(n, h/stride_h, w/stride_w, c), tvm.const(0.0)),
+        name='DilatedInput')
+    return Output

From ac72db01bae70ef51584dcc5baf27e09c350bf87 Mon Sep 17 00:00:00 2001
From: Huyuwei <huyuwei1995@gmail.com>
Date: Sun, 13 Aug 2017 23:28:45 +0800
Subject: [PATCH 2/5] fix pylint

---
 topi/python/topi/nn/dilate.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/topi/python/topi/nn/dilate.py b/topi/python/topi/nn/dilate.py
index c06905c69dba..a75024cd01e7 100644
--- a/topi/python/topi/nn/dilate.py
+++ b/topi/python/topi/nn/dilate.py
@@ -1,8 +1,10 @@
+# pylint: disable=invalid-name
 """Dilation operators"""
 from __future__ import absolute_import as _abs
-from .util import get_const_tuple
 import tvm
 
+from .util import get_const_tuple
+
 @tvm.tag_scope(tag="dilation")
 def dilate_nchw(Input, stride_h, stride_w):
     """Dilate Input with zeros.

From bfc1308a918ad29fed4bbc44b21518b979c14502 Mon Sep 17 00:00:00 2001
From: Huyuwei <huyuwei1995@gmail.com>
Date: Mon, 14 Aug 2017 12:45:02 +0800
Subject: [PATCH 3/5] dilate testcases success

---
 topi/python/topi/nn/dilate.py             | 55 +++++------------------
 topi/python/topi/testing/__init__.py      |  1 +
 topi/python/topi/testing/dilate_python.py | 30 +++++++++++++
 topi/tests/python/test_topi_dilate.py     | 32 +++++++++++++
 4 files changed, 74 insertions(+), 44 deletions(-)
 create mode 100644 topi/python/topi/testing/dilate_python.py
 create mode 100644 topi/tests/python/test_topi_dilate.py

diff --git a/topi/python/topi/nn/dilate.py b/topi/python/topi/nn/dilate.py
index a75024cd01e7..41ff0bfe6db9 100644
--- a/topi/python/topi/nn/dilate.py
+++ b/topi/python/topi/nn/dilate.py
@@ -3,63 +3,30 @@
 from __future__ import absolute_import as _abs
 import tvm
 
-from .util import get_const_tuple
 
 @tvm.tag_scope(tag="dilation")
-def dilate_nchw(Input, stride_h, stride_w):
+def dilate(Input, strides):
     """Dilate Input with zeros.
 
     Parameters
     ----------
     Input : tvm.Tensor
-        Input tensor, layout is NCHW.
+        4-D, can be any layout.
 
-    stride_h : int
-        Dilation extent on H dimension, 1 means no dilation.
-
-    stride_w : int
-        Dilation extent on W dimension, 1 means no dilation.
-
-    Returns
-    -------
-    Output : tvm.Tensor
-        Output tensor, layout is NCHW.
-    """
-    N, C, H, W = get_const_tuple(Input.shape)
-    Output = tvm.compute(
-        (N, C, (H-1)*stride_h+1, (W-1)*stride_w+1),
-        lambda n, c, h, w: tvm.select(
-            tvm.all(tvm.make.EQ(h%stride_h, 0), tvm.make.EQ(w%stride_w, 0)),
-            Input(n, c, h/stride_h, w/stride_w), tvm.const(0.0)),
-        name='DilatedInput')
-    return Output
-
-
-@tvm.tag_scope(tag="dilation")
-def dilate_nhwc(Input, stride_h, stride_w):
-    """Dilate Input with zeros.
-
-    Parameters
-    ----------
-    Input : tvm.Tensor
-        Input tensor, layout is NHWC.
-
-    stride_h : int
-        Dilation extent on H dimension, 1 means no dilation.
-
-    stride_w : int
-        Dilation extent on W dimension, 1 means no dilation.
+    strides : list/tuple of 4 ints
+        Dilation stride on each dimension, 1 means no dilation.
 
     Returns
     -------
     Output : tvm.Tensor
-        Output tensor, layout is NHWC.
+        4-D, the same layout as Input.
     """
-    N, H, W, C = get_const_tuple(Input.shape)
+    A, B, C, D = Input.shape
+    sa, sb, sc, sd = strides
     Output = tvm.compute(
-        (N, (H-1)*stride_h+1, (W-1)*stride_w+1, C),
-        lambda n, h, w, c: tvm.select(
-            tvm.all(tvm.make.EQ(h%stride_h, 0), tvm.make.EQ(w%stride_w, 0)),
-            Input(n, h/stride_h, w/stride_w, c), tvm.const(0.0)),
+        ((A-1)*sa+1, (B-1)*sb+1, (C-1)*sc+1, (D-1)*sd+1),
+        lambda a, b, c, d: tvm.select(
+            tvm.all((a%sa).equal(0), (b%sb).equal(0), (c%sc).equal(0), (d%sd).equal(0)),
+            Input(a/sa, b/sb, c/sc, d/sd), tvm.const(0.0, Input.dtype)),
         name='DilatedInput')
     return Output
diff --git a/topi/python/topi/testing/__init__.py b/topi/python/topi/testing/__init__.py
index 63bc8eb7215a..61fe2a60df91 100644
--- a/topi/python/topi/testing/__init__.py
+++ b/topi/python/topi/testing/__init__.py
@@ -6,3 +6,4 @@
 
 from .conv2d_hwcn_python import conv2d_hwcn_python
 from .conv2d_nchw_python import conv2d_nchw_python
+from .dilate_python import dilate_python
diff --git a/topi/python/topi/testing/dilate_python.py b/topi/python/topi/testing/dilate_python.py
new file mode 100644
index 000000000000..c84d7a737f7a
--- /dev/null
+++ b/topi/python/topi/testing/dilate_python.py
@@ -0,0 +1,30 @@
+# pylint: disable=invalid-name, line-too-long
+"""Dilate operation in python"""
+import numpy as np
+
+
+def dilate_python(input_np, strides):
+    """Dilate operation.
+
+    Parameters
+    ----------
+    input_np : numpy.ndarray
+        4-D, can be any layout.
+
+    strides : list/tuple of 4 ints
+        Dilation stride on each dimension, 1 means no dilation.
+
+    Returns
+    -------
+    output_np : numpy.ndarray
+        4-D, the same layout as Input.
+    """
+    A, B, C, D = input_np.shape
+    sa, sb, sc, sd = strides
+    Ao = (A-1)*sa+1
+    Bo = (B-1)*sb+1
+    Co = (C-1)*sc+1
+    Do = (D-1)*sd+1
+    output_np = np.zeros(shape=(Ao, Bo, Co, Do))
+    output_np[0::sa, 0::sb, 0::sc, 0::sd] = input_np
+    return output_np
diff --git a/topi/tests/python/test_topi_dilate.py b/topi/tests/python/test_topi_dilate.py
new file mode 100644
index 000000000000..79761daf3ac0
--- /dev/null
+++ b/topi/tests/python/test_topi_dilate.py
@@ -0,0 +1,32 @@
+import tvm
+import topi
+import numpy as np
+
+
+def test_dilate():
+    target = 'llvm'
+    ctx = tvm.cpu(0)
+
+    def _test_dilate(input_size, strides):
+        Input = tvm.placeholder((input_size))
+        Output = topi.nn.dilate(Input, strides)
+        schedule = tvm.create_schedule(Output.op)
+        input_np = np.random.uniform(size=input_size).astype(Input.dtype)
+        output_np = topi.testing.dilate_python(input_np, strides)
+        input_tvm = tvm.nd.array(input_np, ctx=ctx)
+        output_tvm = tvm.nd.array(np.zeros(shape=(tvm.ir_pass.Simplify(Output.shape[0]).value,
+            tvm.ir_pass.Simplify(Output.shape[1]).value,
+            tvm.ir_pass.Simplify(Output.shape[2]).value,
+            tvm.ir_pass.Simplify(Output.shape[3]).value)).astype(Output.dtype), ctx=ctx)
+        f = tvm.build(schedule, [Input, Output], target)
+        f(input_tvm, output_tvm)
+        np.testing.assert_allclose(output_tvm.asnumpy(), output_np, rtol=1e-5)
+
+    _test_dilate((1,3,32,32), (1,1,1,1))
+    _test_dilate((1,3,32,32), (1,1,2,2))
+    _test_dilate((1,32,32,3), (1,2,2,1))
+    _test_dilate((32,32,3,1), (2,2,1,1))
+
+
+if __name__ == "__main__":
+    test_dilate()

From b29a0bd942fb39d6d9d47b2b494c37cd69d73f3f Mon Sep 17 00:00:00 2001
From: Huyuwei <huyuwei1995@gmail.com>
Date: Mon, 14 Aug 2017 19:16:15 +0800
Subject: [PATCH 4/5] n-D tensor dilation

---
 topi/python/topi/nn/dilate.py             | 71 +++++++++++++++++++----
 topi/python/topi/testing/dilate_python.py | 35 +++++++----
 topi/tests/python/test_topi_dilate.py     | 19 +++---
 3 files changed, 95 insertions(+), 30 deletions(-)

diff --git a/topi/python/topi/nn/dilate.py b/topi/python/topi/nn/dilate.py
index 41ff0bfe6db9..3363a7f5747e 100644
--- a/topi/python/topi/nn/dilate.py
+++ b/topi/python/topi/nn/dilate.py
@@ -1,4 +1,4 @@
-# pylint: disable=invalid-name
+# pylint: disable=invalid-name, line-too-long
 """Dilation operators"""
 from __future__ import absolute_import as _abs
 import tvm
@@ -11,22 +11,69 @@ def dilate(Input, strides):
     Parameters
     ----------
     Input : tvm.Tensor
-        4-D, can be any layout.
+        n-D, can be any layout.
 
-    strides : list/tuple of 4 ints
+    strides : list / tuple of n ints
         Dilation stride on each dimension, 1 means no dilation.
 
     Returns
     -------
     Output : tvm.Tensor
-        4-D, the same layout as Input.
+        n-D, the same layout as Input.
     """
-    A, B, C, D = Input.shape
-    sa, sb, sc, sd = strides
-    Output = tvm.compute(
-        ((A-1)*sa+1, (B-1)*sb+1, (C-1)*sc+1, (D-1)*sd+1),
-        lambda a, b, c, d: tvm.select(
-            tvm.all((a%sa).equal(0), (b%sb).equal(0), (c%sc).equal(0), (d%sd).equal(0)),
-            Input(a/sa, b/sb, c/sc, d/sd), tvm.const(0.0, Input.dtype)),
-        name='DilatedInput')
+    n = len(Input.shape)
+    assert n <= 5, \
+        "Dimension of input tensor cannot exceed 5"
+    assert len(strides) == n, \
+        "Input dimension and strides size dismatch : %d vs %d" %(n, len(strides))
+    output_size = ()
+    for i in range(n):
+        output_size += (tvm.ir_pass.Simplify((Input.shape[i]-1)*strides[i]+1),)
+
+    if n == 5:
+        Output = tvm.compute(
+            (output_size),
+            lambda *i: tvm.select(
+                tvm.all((i[0]%strides[0]).equal(0),
+                        (i[1]%strides[1]).equal(0),
+                        (i[2]%strides[2]).equal(0),
+                        (i[3]%strides[3]).equal(0),
+                        (i[4]%strides[4]).equal(0)),
+                Input(i[0]/strides[0], i[1]/strides[1], i[2]/strides[2], i[3]/strides[3], i[4]/strides[4]),
+                tvm.const(0.0, Input.dtype)), name='DilatedInput')
+    elif n == 4:
+        Output = tvm.compute(
+            (output_size),
+            lambda *i: tvm.select(
+                tvm.all((i[0]%strides[0]).equal(0),
+                        (i[1]%strides[1]).equal(0),
+                        (i[2]%strides[2]).equal(0),
+                        (i[3]%strides[3]).equal(0)),
+                Input(i[0]/strides[0], i[1]/strides[1], i[2]/strides[2], i[3]/strides[3]),
+                tvm.const(0.0, Input.dtype)), name='DilatedInput')
+    elif n == 3:
+        Output = tvm.compute(
+            (output_size),
+            lambda *i: tvm.select(
+                tvm.all((i[0]%strides[0]).equal(0),
+                        (i[1]%strides[1]).equal(0),
+                        (i[2]%strides[2]).equal(0)),
+                Input(i[0]/strides[0], i[1]/strides[1], i[2]/strides[2]),
+                tvm.const(0.0, Input.dtype)), name='DilatedInput')
+    elif n == 2:
+        Output = tvm.compute(
+            (output_size),
+            lambda *i: tvm.select(
+                tvm.all((i[0]%strides[0]).equal(0),
+                        (i[1]%strides[1]).equal(0)),
+                Input(i[0]/strides[0], i[1]/strides[1]),
+                tvm.const(0.0, Input.dtype)), name='DilatedInput')
+    else: # n == 1
+        Output = tvm.compute(
+            (output_size),
+            lambda *i: tvm.select(
+                tvm.all((i[0]%strides[0]).equal(0)),
+                Input(i[0]/strides[0]),
+                tvm.const(0.0, Input.dtype)), name='DilatedInput')
+
     return Output
diff --git a/topi/python/topi/testing/dilate_python.py b/topi/python/topi/testing/dilate_python.py
index c84d7a737f7a..68fec6ae9608 100644
--- a/topi/python/topi/testing/dilate_python.py
+++ b/topi/python/topi/testing/dilate_python.py
@@ -9,22 +9,35 @@ def dilate_python(input_np, strides):
     Parameters
     ----------
     input_np : numpy.ndarray
-        4-D, can be any layout.
+        n-D, can be any layout.
 
-    strides : list/tuple of 4 ints
+    strides : list / tuple of n ints
         Dilation stride on each dimension, 1 means no dilation.
 
     Returns
     -------
     output_np : numpy.ndarray
-        4-D, the same layout as Input.
+        n-D, the same layout as Input.
     """
-    A, B, C, D = input_np.shape
-    sa, sb, sc, sd = strides
-    Ao = (A-1)*sa+1
-    Bo = (B-1)*sb+1
-    Co = (C-1)*sc+1
-    Do = (D-1)*sd+1
-    output_np = np.zeros(shape=(Ao, Bo, Co, Do))
-    output_np[0::sa, 0::sb, 0::sc, 0::sd] = input_np
+    n = len(input_np.shape)
+    assert n <= 5, \
+        "Dimension of input tensor cannot exceed 5"
+    assert len(strides) == n, \
+        "Input dimension and strides size dismatch : %d vs %d" %(n, len(strides))
+    output_size = ()
+    for i in range(n):
+        output_size += ((input_np.shape[i]-1)*strides[i]+1,)
+    output_np = np.zeros(shape=output_size)
+
+    if n == 5:
+        output_np[0::strides[0], 0::strides[1], 0::strides[2], 0::strides[3], 0::strides[4]] = input_np
+    elif n == 4:
+        output_np[0::strides[0], 0::strides[1], 0::strides[2], 0::strides[3]] = input_np
+    elif n == 3:
+        output_np[0::strides[0], 0::strides[1], 0::strides[2]] = input_np
+    elif n == 2:
+        output_np[0::strides[0], 0::strides[1]] = input_np
+    else: # n == 1
+        output_np[0::strides[0]] = input_np
+
     return output_np
diff --git a/topi/tests/python/test_topi_dilate.py b/topi/tests/python/test_topi_dilate.py
index 79761daf3ac0..eeb4114de42f 100644
--- a/topi/tests/python/test_topi_dilate.py
+++ b/topi/tests/python/test_topi_dilate.py
@@ -14,18 +14,23 @@ def _test_dilate(input_size, strides):
         input_np = np.random.uniform(size=input_size).astype(Input.dtype)
         output_np = topi.testing.dilate_python(input_np, strides)
         input_tvm = tvm.nd.array(input_np, ctx=ctx)
-        output_tvm = tvm.nd.array(np.zeros(shape=(tvm.ir_pass.Simplify(Output.shape[0]).value,
-            tvm.ir_pass.Simplify(Output.shape[1]).value,
-            tvm.ir_pass.Simplify(Output.shape[2]).value,
-            tvm.ir_pass.Simplify(Output.shape[3]).value)).astype(Output.dtype), ctx=ctx)
+        output_size = ()
+        for i in range(len(input_size)):
+            output_size += (tvm.ir_pass.Simplify(Output.shape[i]).value,)
+        output_tvm = tvm.nd.array(np.zeros(shape=output_size).astype(Output.dtype), ctx=ctx)
         f = tvm.build(schedule, [Input, Output], target)
         f(input_tvm, output_tvm)
         np.testing.assert_allclose(output_tvm.asnumpy(), output_np, rtol=1e-5)
 
+
+    _test_dilate((32,), (2,))
+    _test_dilate((32,32), (2,2))
+    _test_dilate((3,32,32), (1,2,2))
+    _test_dilate((32,32,3), (2,2,1))
     _test_dilate((1,3,32,32), (1,1,1,1))
-    _test_dilate((1,3,32,32), (1,1,2,2))
-    _test_dilate((1,32,32,3), (1,2,2,1))
-    _test_dilate((32,32,3,1), (2,2,1,1))
+    _test_dilate((1,3,32,32), (2,2,2,2))
+    _test_dilate((1,32,32,3,3), (1,1,1,1,1))
+    _test_dilate((1,32,32,3,3), (2,2,2,2,2))
 
 
 if __name__ == "__main__":

From 7a2daef147b0773a679e18ce16c41566c7e657aa Mon Sep 17 00:00:00 2001
From: Huyuwei <huyuwei1995@gmail.com>
Date: Tue, 15 Aug 2017 00:56:07 +0800
Subject: [PATCH 5/5] support arbitrary dimension

---
 topi/python/topi/nn/dilate.py             | 61 +++++------------------
 topi/python/topi/testing/dilate_python.py | 18 ++-----
 topi/tests/python/test_topi_dilate.py     |  5 +-
 3 files changed, 19 insertions(+), 65 deletions(-)

diff --git a/topi/python/topi/nn/dilate.py b/topi/python/topi/nn/dilate.py
index 3363a7f5747e..2b3b2f424aa6 100644
--- a/topi/python/topi/nn/dilate.py
+++ b/topi/python/topi/nn/dilate.py
@@ -1,4 +1,4 @@
-# pylint: disable=invalid-name, line-too-long
+# pylint: disable=invalid-name
 """Dilation operators"""
 from __future__ import absolute_import as _abs
 import tvm
@@ -22,58 +22,23 @@ def dilate(Input, strides):
         n-D, the same layout as Input.
     """
     n = len(Input.shape)
-    assert n <= 5, \
-        "Dimension of input tensor cannot exceed 5"
     assert len(strides) == n, \
         "Input dimension and strides size dismatch : %d vs %d" %(n, len(strides))
     output_size = ()
     for i in range(n):
         output_size += (tvm.ir_pass.Simplify((Input.shape[i]-1)*strides[i]+1),)
 
-    if n == 5:
-        Output = tvm.compute(
-            (output_size),
-            lambda *i: tvm.select(
-                tvm.all((i[0]%strides[0]).equal(0),
-                        (i[1]%strides[1]).equal(0),
-                        (i[2]%strides[2]).equal(0),
-                        (i[3]%strides[3]).equal(0),
-                        (i[4]%strides[4]).equal(0)),
-                Input(i[0]/strides[0], i[1]/strides[1], i[2]/strides[2], i[3]/strides[3], i[4]/strides[4]),
-                tvm.const(0.0, Input.dtype)), name='DilatedInput')
-    elif n == 4:
-        Output = tvm.compute(
-            (output_size),
-            lambda *i: tvm.select(
-                tvm.all((i[0]%strides[0]).equal(0),
-                        (i[1]%strides[1]).equal(0),
-                        (i[2]%strides[2]).equal(0),
-                        (i[3]%strides[3]).equal(0)),
-                Input(i[0]/strides[0], i[1]/strides[1], i[2]/strides[2], i[3]/strides[3]),
-                tvm.const(0.0, Input.dtype)), name='DilatedInput')
-    elif n == 3:
-        Output = tvm.compute(
-            (output_size),
-            lambda *i: tvm.select(
-                tvm.all((i[0]%strides[0]).equal(0),
-                        (i[1]%strides[1]).equal(0),
-                        (i[2]%strides[2]).equal(0)),
-                Input(i[0]/strides[0], i[1]/strides[1], i[2]/strides[2]),
-                tvm.const(0.0, Input.dtype)), name='DilatedInput')
-    elif n == 2:
-        Output = tvm.compute(
-            (output_size),
-            lambda *i: tvm.select(
-                tvm.all((i[0]%strides[0]).equal(0),
-                        (i[1]%strides[1]).equal(0)),
-                Input(i[0]/strides[0], i[1]/strides[1]),
-                tvm.const(0.0, Input.dtype)), name='DilatedInput')
-    else: # n == 1
-        Output = tvm.compute(
-            (output_size),
-            lambda *i: tvm.select(
-                tvm.all((i[0]%strides[0]).equal(0)),
-                Input(i[0]/strides[0]),
-                tvm.const(0.0, Input.dtype)), name='DilatedInput')
+    def _dilate(data, *indices):
+        not_zero = (indices[0]%strides[0]).equal(0)
+        index_tuple = ()
+        for i in range(n):
+            index_tuple += (indices[i]/strides[i],)
+            not_zero = tvm.all(not_zero, (indices[i]%strides[i]).equal(0))
+        return tvm.select(not_zero, data[index_tuple], tvm.const(0.0, data.dtype))
+
+    Output = tvm.compute(
+        (output_size),
+        lambda *indices: _dilate(Input, *indices),
+        name='DilatedInput')
 
     return Output
diff --git a/topi/python/topi/testing/dilate_python.py b/topi/python/topi/testing/dilate_python.py
index 68fec6ae9608..89ac2c109fb6 100644
--- a/topi/python/topi/testing/dilate_python.py
+++ b/topi/python/topi/testing/dilate_python.py
@@ -1,4 +1,4 @@
-# pylint: disable=invalid-name, line-too-long
+# pylint: disable=invalid-name
 """Dilate operation in python"""
 import numpy as np
 
@@ -20,24 +20,14 @@ def dilate_python(input_np, strides):
         n-D, the same layout as Input.
     """
     n = len(input_np.shape)
-    assert n <= 5, \
-        "Dimension of input tensor cannot exceed 5"
     assert len(strides) == n, \
         "Input dimension and strides size dismatch : %d vs %d" %(n, len(strides))
     output_size = ()
+    no_zero = ()
     for i in range(n):
         output_size += ((input_np.shape[i]-1)*strides[i]+1,)
+        no_zero += ((range(0, output_size[i], strides[i])),)
     output_np = np.zeros(shape=output_size)
-
-    if n == 5:
-        output_np[0::strides[0], 0::strides[1], 0::strides[2], 0::strides[3], 0::strides[4]] = input_np
-    elif n == 4:
-        output_np[0::strides[0], 0::strides[1], 0::strides[2], 0::strides[3]] = input_np
-    elif n == 3:
-        output_np[0::strides[0], 0::strides[1], 0::strides[2]] = input_np
-    elif n == 2:
-        output_np[0::strides[0], 0::strides[1]] = input_np
-    else: # n == 1
-        output_np[0::strides[0]] = input_np
+    output_np[np.ix_(*no_zero)] = input_np
 
     return output_np
diff --git a/topi/tests/python/test_topi_dilate.py b/topi/tests/python/test_topi_dilate.py
index eeb4114de42f..0d2014535ca4 100644
--- a/topi/tests/python/test_topi_dilate.py
+++ b/topi/tests/python/test_topi_dilate.py
@@ -22,15 +22,14 @@ def _test_dilate(input_size, strides):
         f(input_tvm, output_tvm)
         np.testing.assert_allclose(output_tvm.asnumpy(), output_np, rtol=1e-5)
 
-
     _test_dilate((32,), (2,))
     _test_dilate((32,32), (2,2))
-    _test_dilate((3,32,32), (1,2,2))
-    _test_dilate((32,32,3), (2,2,1))
     _test_dilate((1,3,32,32), (1,1,1,1))
     _test_dilate((1,3,32,32), (2,2,2,2))
     _test_dilate((1,32,32,3,3), (1,1,1,1,1))
     _test_dilate((1,32,32,3,3), (2,2,2,2,2))
+    _test_dilate((1,32,32,32,3,3), (1,1,1,2,2,2))
+    _test_dilate((1,32,32,32,3,3), (2,2,2,1,1,1))
 
 
 if __name__ == "__main__":