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,