<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def MatrixDiagPart(a):
 """ Batched diag op that returns only the diagonal elements. """

    r = np.zeros(a.shape[:-2] + (min(a.shape[-2:]),))
    for coord in np.ndindex(a.shape[:-2]):
        pos = coord + (Ellipsis,)
        r[pos] = np.diagonal(a[pos])
    return r,