fix prelu, now can use on 2d input and add one test (apache#2875)

MarisaKirisame · Apr 9, 2019 · 2fc7662 · 2fc7662
1 parent 997ae1e
commit 2fc7662
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 2 deletions.
diff --git a/topi/include/topi/nn.h b/topi/include/topi/nn.h
@@ -97,7 +97,6 @@ inline tvm::Tensor prelu(const tvm::Tensor &x,
                          const int axis = 1,
                          std::string name = "tensor",
                          std::string tag = kBroadcast) {
-  CHECK_EQ(4, x->shape.size());
   CHECK((size_t)axis < x->shape.size()) <<
         "Wrong axis ("  << axis << ")value. ";
   CHECK(topi::detail::GetConstInt(slope->shape[0]) ==

diff --git a/topi/python/topi/nn/elemwise.py b/topi/python/topi/nn/elemwise.py
@@ -69,7 +69,7 @@ def prelu(x, slope, axis=1):
         [http://arxiv.org/pdf/1502.01852v1.pdf]
     """
 
-    assert len(x.shape) == 4 and len(slope.shape) == 1
+    assert len(slope.shape) == 1
     assert axis < len(x.shape)
     assert get_const_int(slope.shape[0]) == get_const_int(x.shape[axis])
 

diff --git a/topi/tests/python/test_topi_relu.py b/topi/tests/python/test_topi_relu.py
@@ -83,6 +83,7 @@ def test_leaky_relu():
 def test_prelu():
     verify_prelu((1, 3, 2, 2), (3,), 1, (3, 1, 1))
     verify_prelu((1, 3, 2, 2), (2,), 2, (2, 1))
+    verify_prelu((1, 3), (3,), 1, (3, ))
 
 if __name__ == "__main__":
     test_schedule_big_array()