diff --git a/python/mxnet/numpy/multiarray.py b/python/mxnet/numpy/multiarray.py index 39c3bb6bb761..ce3f853b041b 100644 --- a/python/mxnet/numpy/multiarray.py +++ b/python/mxnet/numpy/multiarray.py @@ -7517,10 +7517,10 @@ def where(condition, x=None, y=None): @set_module('mxnet.numpy') def diagonal(a, offset=0, axis1=0, axis2=1): """ - If a is 2-D, returns the diagonal of a with the given offset, i.e., the collection of elements of - the form a[i, i+offset]. If a has more than two dimensions, then the axes specified by axis1 and - axis2 are used to determine the 2-D sub-array whose diagonal is returned. The shape of the - resulting array can be determined by removing axis1 and axis2 and appending an index to the + If a is 2-D, returns the diagonal of a with the given offset, i.e., the collection of elements of + the form a[i, i+offset]. If a has more than two dimensions, then the axes specified by axis1 and + axis2 are used to determine the 2-D sub-array whose diagonal is returned. The shape of the + resulting array can be determined by removing axis1 and axis2 and appending an index to the right equal to the size of the resulting diagonals. Parameters @@ -7560,9 +7560,8 @@ def diagonal(a, offset=0, axis1=0, axis2=1): [2, 3]], [[4, 5], [6, 7]]]) - >>> np.diagonal(a, 0, 0, 1) + >>> np.diagonal(a, 0, 0, 1) array([[0, 6], [1, 7]]) """ return _mx_nd_np.diagonal(a, offset=offset, axis1=axis1, axis2=axis2) - diff --git a/src/operator/numpy/np_matrix_op-inl.h b/src/operator/numpy/np_matrix_op-inl.h index 41c60442e763..38af423b3d3a 100644 --- a/src/operator/numpy/np_matrix_op-inl.h +++ b/src/operator/numpy/np_matrix_op-inl.h @@ -1316,8 +1316,8 @@ void NumpyDiagonalOpImpl(const TBlob& in_data, stride1 + stride2, offset, oshape[odim - 1]); } else { Kernel, xpu>::Launch(s, dsize, out_data.dptr(), - in_data.dptr(), Shape3(oleading, obody, otrailing), Shape3(ileading, ibody, itrailing), - stride1 + stride2, offset, oshape[odim - 1]); + in_data.dptr(), Shape3(oleading, obody, otrailing), + Shape3(ileading, ibody, itrailing), stride1 + stride2, offset, oshape[odim - 1]); } }); }); @@ -1342,7 +1342,8 @@ void NumpyDiagonalOpForward(const nnvm::NodeAttrs& attrs, const mxnet::TShape& oshape = outputs[0].shape_; const NumpyDiagonalParam& param = nnvm::get(attrs.parsed); - NumpyDiagonalOpImpl(in_data, out_data, ishape, oshape, out_data.Size(), param, s, req); + NumpyDiagonalOpImpl(in_data, out_data, ishape, oshape, + out_data.Size(), param, s, req); } template @@ -1363,7 +1364,8 @@ void NumpyDiagonalOpBackward(const nnvm::NodeAttrs& attrs, const mxnet::TShape& oshape = outputs[0].shape_; const NumpyDiagonalParam& param = nnvm::get(attrs.parsed); - NumpyDiagonalOpImpl(in_data, out_data, oshape, ishape, in_data.Size(), param, s, req); + NumpyDiagonalOpImpl(in_data, out_data, oshape, ishape, + in_data.Size(), param, s, req); } struct NumpyDiagflatParam : public dmlc::Parameter {