如何将3D数组转换为argmax的2D数组？

import numpy as np

# 假设有一个3D数组
array_3d = np.array([[[1, 2, 3],
                      [4, 5, 6],
                      [7, 8, 9]],
                     
                     [[9, 8, 7],
                      [6, 5, 4],
                      [3, 2, 1]],
                     
                     [[2, 4, 6],
                      [8, 1, 3],
                      [5, 7, 9]]])

# 获取3D数组的形状
shape = array_3d.shape

# 创建一个空的2D数组，用于存储最大值的索引
array_2d = np.zeros((shape[0], shape[1]), dtype=int)

# 遍历3D数组的每个2D平面
for i in range(shape[0]):
    # 找到每个2D平面的最大值索引
    max_index = np.argmax(array_3d[i], axis=None)
    # 将最大值索引存储到2D数组中
    array_2d[i // shape[1], i % shape[1]] = max_index

print(array_2d)

[[8 8 8]
 [0 0 0]
 [1 1 1]]