code: def MatrixDiagPart(a):
      
      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,
docs: Batched diag op that returns only the diagonal elements.