[BUGFIX] Minor type issues in Squeeze (#16448)

* minor * test * Retrigger * Retrigger
apache · Oct 13, 2019 · 1e8cc90 · 1e8cc90
1 parent 4dee4ee
commit 1e8cc90
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 10 deletions.
diff --git a/src/operator/numpy/np_matrix_op.cc b/src/operator/numpy/np_matrix_op.cc
@@ -248,7 +248,7 @@ NNVM_REGISTER_OP(_np_squeeze)
 .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>)
 .set_attr<FCompute>("FCompute<cpu>", UnaryOp::IdentityCompute<cpu>)
 .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_squeeze"})
-.add_argument("a", "NDArray-or-Symbol[]", "data to squeeze")
+.add_argument("a", "NDArray-or-Symbol", "data to squeeze")
 .add_arguments(SqueezeParam::__FIELDS__());
 
 bool ConcatShape(const nnvm::NodeAttrs& attrs,

diff --git a/src/operator/tensor/matrix_op.cc b/src/operator/tensor/matrix_op.cc
@@ -1003,7 +1003,7 @@ Examples::
 .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>)
 .set_attr<FCompute>("FCompute<cpu>", UnaryOp::IdentityCompute<cpu>)
 .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_squeeze"})
-.add_argument("data", "NDArray-or-Symbol[]", "data to squeeze")
+.add_argument("data", "NDArray-or-Symbol", "data to squeeze")
 .add_arguments(SqueezeParam::__FIELDS__());
 
 NNVM_REGISTER_OP(_backward_squeeze)

diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py
@@ -228,11 +228,11 @@ def __init__(self):
 
         def hybrid_forward(self, F, x1, x2):
             return F.np.ldexp(x1, x2)
-        
+
     def _np_ldexp(x1, x2):
         return x1 * _np.power(2.0, x2)
 
-    def dldx(x1, x2): 
+    def dldx(x1, x2):
         grad_a = _np.power(2.0, x2)
         grad_b = _np_ldexp(x1, x2) * _np.log(2.0)
         if len(x1) == 1:
@@ -241,7 +241,7 @@ def dldx(x1, x2):
             grad_b = _np.sum(grad_b)
         return [grad_a, grad_b]
 
-    shapes = [ 
+    shapes = [
         ((3, 1), (3, 1)),
         ((3, 1, 2), (3, 1, 2)),
         ((1, ),(1, )),
@@ -250,15 +250,15 @@ def dldx(x1, x2):
         ((3, 0), (3, 0)),  # zero-size shape
         ((0, 1), (0, 1)),  # zero-size shape
         ((2, 0, 2), (2, 0, 2)),  # zero-size shape
-        ] 
+        ]
 
     for hybridize in [True, False]:
         for shape1, shape2 in shapes:
             for dtype in [_np.float16, _np.float32, _np.float64]:
                 test_ldexp = TestLdexp()
                 if hybridize:
                     test_ldexp.hybridize()
-                x1 = rand_ndarray(shape=shape1, dtype=dtype).as_np_ndarray() 
+                x1 = rand_ndarray(shape=shape1, dtype=dtype).as_np_ndarray()
                 x1.attach_grad()
                 x2 = rand_ndarray(shape=shape2, dtype=dtype).as_np_ndarray()
                 x2.attach_grad()
@@ -997,13 +997,13 @@ def __init__(self, axis):
             self._axis = axis
 
         def hybrid_forward(self, F, x):
-            return F.np.squeeze(x, axis=self._axis)
+            return F.np.squeeze(x, self._axis)
 
     for shape, axis in config:
         data_np = _np.random.uniform(size=shape)
         data_mx = np.array(data_np, dtype=data_np.dtype)
-        ret_np = _np.squeeze(data_np, axis=axis)
-        ret_mx = np.squeeze(data_mx, axis=axis)
+        ret_np = _np.squeeze(data_np, axis)
+        ret_mx = np.squeeze(data_mx, axis)
         assert_almost_equal(ret_mx.asnumpy(), ret_np, rtol=1e-5, atol=1e-6, use_broadcast=False)
 
         net = TestSqueeze(axis)