Как сформировать индекс массива 3D numpy

Question

У меня есть трехмерный массив numpy, b, с shape=(3,2,3)

В принципе, я использовал argmax(), чтобы найти максимальные индексы первой оси, и мне нужно восстановить значения из индексов

Вот мой тестовый код,

#Test array, reshape, find max indices
a=np.array([[1, 2, 3], [4, 5, 6], [13, 14, 15], [10, 17, 12], [7, 8, 9], [16, 11, 18]])
b=np.reshape(a, (3,2,3))
max_indices = np.argmax(a=b, axis=0)


#I need to form an index that finds the following entries
print(b[1,0,0])  # 13
print(b[1,0,1])  # 14
print(b[1,0,2])  # 15
print(b[2,1,0])  # 16
print(b[1,1,1])  # 17
print(b[2,1,2])  # 18

#I tried the following,
m, n = b.shape[:2]
I, J = np.ogrid[:m, :n]

max_values = b[max_indices, J, I]

Но есть ошибка трансляции.....

Любая помощь с благодарностью принята.

Дуг