Взять конкретный 2d массив из 3d в numpy

Question

Есть ли способ избежать использования цикла for и получить результат, просто вызвав arr с некоторой индексацией?Потенциально dim1 будет равно 50 000, dim2 - 1000, dim3 - 3.

import numpy as np
dim1 = 10
dim2 = 2
dim3 = 3
arr = np.arange(60).reshape(dim1,dim2,dim3)
arr2 = np.arange(dim1*dim2).reshape(dim1,dim2)
np.mod(arr2,dim3,out=arr2)

res = []
rng = np.arange(dim1)
for x in range(dim2):
    sl = arr2[:,x]
    temp = arr[rng,x,sl]
    res.append(temp)
res = np.asarray(res).T

По сути, я хотел бы извлечь значения из arr, который является трехмерным массивом, однако матрица arr2 указываеткакие столбцы выбрать.

Best