Tengo una gran variedad de formas arbitrarias, por ejemplo:

a = array([[[ 1,  2],
            [ 3,  4],
            [ 8,  6]],

          [[ 7,  8],
           [ 9,  8],
           [ 3, 12]]])
a.shape = (2, 3, 2)

y un resultado de argmax sobre el último eje:

np.argmax(a, axis=-1) = array([[1, 1, 0],
                               [1, 0, 1]])

Me gustaría obtener max:

np.max(a, axis=-1) = array([[ 2,  4,  8],
                            [ 8,  9, 12]])

Pero sin recalcular todo. He intentado:

a[np.arange(len(a)), np.argmax(a, axis=-1)]

Pero tengo:

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (2,) (2,3) 

¿Cómo hacerlo? Pregunta similar para 2-d:numpy 2d array max / argmax