def index_at(idx, shape, axis=-1):
    if axis < 0:
        axis += len(shape)
shape = shape[:axis] + shape[axis + 1:]
index = list(np.ix_(*[np.arange(n) for n in shape]))
index.insert(axis, idx)
return tuple(index)
a = np.random.randint(0, 10, (3, 4, 5))
axis = 1
idx = np.argmax(a, axis=axis)
print(a[index_at(idx, a.shape, axis=axis)])
print(np.max(a, axis=axis))
преобразование numpy из 3d в 2d на основе массива 2d-масок