From 593b29cf41b8c59ed9720b72e64ae4b99ef1c185 Mon Sep 17 00:00:00 2001 From: Fan Date: Tue, 2 Jul 2019 00:35:32 +0800 Subject: [PATCH] add doc for trace --- python/mxnet/_numpy_op_doc.py | 49 +++++++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/python/mxnet/_numpy_op_doc.py b/python/mxnet/_numpy_op_doc.py index 5543ebc8e8c9..9fd95e6c4317 100644 --- a/python/mxnet/_numpy_op_doc.py +++ b/python/mxnet/_numpy_op_doc.py @@ -52,3 +52,52 @@ def _np_zeros_like(a): Array of zeros with the same shape and type as `a`. """ pass + + +def _np_trace(a, offset=0, axis1=0, axis2=1, out=None): + """trace(a, offset=0, axis1=0, axis2=1, out=None) + + Return the sum along diagonals of the array. + + If `a` is 2-D, the sum along its diagonal with the given offset + is returned, i.e., the sum of elements ``a[i,i+offset]`` for all i. + + If `a` has more than two dimensions, then the axes specified by axis1 and + axis2 are used to determine the 2-D sub-arrays whose traces are returned. + The shape of the resulting array is the same as that of `a` with `axis1` + and `axis2` removed. + + Parameters + ---------- + a : ndarray + Input array, from which the diagonals are taken. + offset : int, optional + Offset of the diagonal from the main diagonal. Can be both positive + and negative. Defaults to 0. + axis1, axis2 : int, optional + Axes to be used as the first and second axis of the 2-D sub-arrays + from which the diagonals should be taken. Defaults are the first two + axes of `a`. + out : ndarray, optional + Array into which the output is placed. It must be of the right shape + and right type to hold the output. + + Returns + ------- + sum_along_diagonals : ndarray + If `a` is 2-D, the sum along the diagonal is returned. If `a` has + larger dimensions, then an array of sums along diagonals is returned. + + Examples + -------- + >>> a = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) + >>> np.trace(a) + array(3.) + >>> a = np.arange(8).reshape((2, 2, 2)) + >>> np.trace(a) + array([6., 8.]) + >>> a = np.arange(24).reshape((2, 2, 2, 3)) + >>> np.trace(a).shape + (2, 3) + """ + pass