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

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]) ==

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])
 

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()