change unittest and fix

apache · Nov 15, 2019 · 9006455 · 9006455
1 parent e6894e1
commit 9006455
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/src/operator/numpy/np_matrix_op-inl.h b/src/operator/numpy/np_matrix_op-inl.h
@@ -1099,7 +1099,7 @@ void NumpyDiagOpImpl(const TBlob &in_data,
       MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
         Kernel<diag_gen<req_type, back>, xpu>::Launch(
             s, dsize, out_data.dptr<DType>(), in_data.dptr<DType>(),
-            Shape2(oshape[0], oshape[1]), param.k);
+            Shape2(oshape[0], oshape[1]), k);
       });
     });
   }

diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py
@@ -4499,7 +4499,7 @@ def __init__(self, k=0):
         def hybrid_forward(self, F, a):
             return F.np.diag(a, k=self._k)
 
-    shapes = [(2,), 5, (1, 5), (2, 2), (2, 5), (3, 3), (4, 3)]
+    shapes = [(), (2,), (1, 5), (2, 2), (2, 5), (3, 3), (4, 3)]
     dtypes = [np.int8, np.uint8, np.int32, np.int64, np.float16, np.float32, np.float64]
     range_k = 6
     combination = itertools.product([False, True], shapes, dtypes, list(range(-range_k, range_k)))
@@ -4521,10 +4521,9 @@ def hybrid_forward(self, F, a):
 
         # check backward function 
         mx_out.backward()
-        np_backward = 0
-        if type(shape) == int:
-            np_backward = np.ones(shape)
-        elif len(shape) < 2:
+        if len(shape) == 0:
+            np_backward = np.array(())
+        elif len(shape) == 1:
             np_backward = np.ones(shape[0])
         else:
             np_backward = np.zeros(shape)