如何在numpy中找到二维数组的argmax？

import numpy as np

# 创建一个二维数组
arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# 找到整个数组的最大值所在的索引
max_index = np.argmax(arr)
print("整个数组的最大值所在的索引：", max_index)

# 沿着行的方向找到每行的最大值所在的索引
row_max_index = np.argmax(arr, axis=1)
print("每行的最大值所在的索引：", row_max_index)

# 沿着列的方向找到每列的最大值所在的索引
col_max_index = np.argmax(arr, axis=0)
print("每列的最大值所在的索引：", col_max_index)

整个数组的最大值所在的索引： 8
每行的最大值所在的索引： [2 2 2]
每列的最大值所在的索引： [2 2 2]