Qual é a maneira mais elegante de acessar uma matriz dimensional n com uma matriz dimensional (n-1) ao longo de uma determinada dimensão, como no exemplo fictício

a = np.random.random_sample((3,4,4))
b = np.random.random_sample((3,4,4))
idx = np.argmax(a, axis=0)

Como posso acessar agora comidx a para obter o máximo ema como se eu tivesse usadoa.max(axis=0)? ou como recuperar os valores especificados poridx nob?

Pensei em usarnp.meshgrid mas acho que é um exagero. Observe que a dimensãoaxis pode ser qualquer eixo útil (0,1,2) e não é conhecido antecipadamente. Existe uma maneira elegante de fazer isso?