Trace of each matrix in a numpy array

Question

Given a list of multiple matrices, I want to compute the trace of each matrix.

np.asarray([[[1,2,3],
             [4,5,6],
             [7,8,9]],
           [[10,11,12],
            [13,14,15],
            [16,17,18]]])

The result should be:

np.asarray([15, 42])

Answer 1

Use axis1=1 and axis2=2 arguments in np.trace.

np.trace(x, axis1=1, axis2=2)

Answer 2

Here's a way using advanced indexing and sum along the second axis:

a = np.asarray([[[1,2,3],
                 [4,5,6],
                 [7,8,9]],
               [[10,11,12],
                [13,14,15],
                [16,17,18]]])

a[:, range(a.shape[1]), range(a.shape[2])].sum(1)
# array([15, 42])