Squeeze bug fix. (apache#506)

tqchen · May 29, 2018 · 510cd5e · 510cd5e
1 parent 69d5fca
commit 510cd5e
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
diff --git a/nnvm/src/top/tensor/transform.cc b/nnvm/src/top/tensor/transform.cc
@@ -638,11 +638,14 @@ inline bool SqueezeShape(const nnvm::NodeAttrs& attrs,
   } else {
     std::unordered_set<dim_t> axis_checker;
     for (size_t i = 0; i < param.axis.ndim(); ++i) {
+      int real_axis;
       if (param.axis[i] < 0) {
-        int real_axis = param.axis[i] + static_cast<int>(shp.ndim());
-        CHECK(real_axis < static_cast<int>(shp.ndim()) && real_axis >= 0);
-        axis_checker.insert(real_axis);
+        real_axis = param.axis[i] + static_cast<int>(shp.ndim());
+      } else {
+        real_axis = param.axis[i];
       }
+      CHECK(real_axis < static_cast<int>(shp.ndim()) && real_axis >= 0);
+      axis_checker.insert(real_axis);
     }
     for (size_t i = 0; i < shp.ndim(); ++i) {
       if (axis_checker.find(i) == axis_checker.end()) {

diff --git a/nnvm/tests/python/unittest/test_infer_shape.py b/nnvm/tests/python/unittest/test_infer_shape.py
@@ -116,6 +116,24 @@ def test_flatten():
     sdict = infer_shape(y)
     assert(sdict["y"][0] == [10, 200])
 
+def test_squeeze():
+    x = sym.Variable("x", shape=(1, 1, 1, 10))
+    y = sym.squeeze(x, axis=(1,2), name='squeeze')
+    sdict = infer_shape(y)
+    assert(sdict['squeeze'][0] == [1, 10])
+
+    x = sym.Variable("x", shape=(1, 3, 1))
+    y = sym.squeeze(x, name='squeeze')
+    sdict = infer_shape(y)
+    assert(sdict['squeeze'][0] == [3])
+
+    y = sym.squeeze(x, axis=(0), name='squeeze')
+    sdict = infer_shape(y)
+    assert(sdict['squeeze'][0] == [3, 1])
+
+    y = sym.squeeze(x, axis=(0,2), name='squeeze')
+    sdict = infer_shape(y)
+    assert(sdict['squeeze'][0] == [3])
 
 # Level 2
 def test_conv2d():
@@ -331,3 +349,4 @@ def check(in_shape, out_shape, **kwargs):
     test_reduce()
     test_transpose()
     test_prelu()
+    test_squeeze()