Как перебрать срезы в последнем измерении

Question

Например, у вас есть массив

a = np.array([[[ 0,  1,  2],
               [ 3,  4,  5]],

              [[ 6,  7,  8],
               [ 9, 10, 11]]])

Мы хотим перебрать срезы в последнем измерении, т.е. [0,1,2], [3,4,5], [6,7 , 8], [9,10,11]. Любой способ достичь этого без для l oop? Спасибо!

Пробовал это, но это не работает, потому что numpy не интерпретирует кортеж так, как мы хотели - a[(0, 0),:] не то же самое, что a[0, 0, :]

[a[i,:] for i in zip(*product(*(range(ii) for ii in a.shape[:-1])))]

В целом, есть ли способ для последних k измерений? Что-то эквивалентное циклу a[i,j,k, ...].

Answer 1

In [26]: a = np.array([[[ 0,  1,  2], 
    ...:                [ 3,  4,  5]], 
    ...:  
    ...:               [[ 6,  7,  8], 
    ...:                [ 9, 10, 11]]])                                                                
In [27]: [a[i,j,:] for i in range(2) for j in range(2)]                                                
Out[27]: [array([0, 1, 2]), array([3, 4, 5]), array([6, 7, 8]), array([ 9, 10, 11])]

или

In [31]: list(np.ndindex(2,2))                                                                         
Out[31]: [(0, 0), (0, 1), (1, 0), (1, 1)]
In [32]: [a[i,j] for i,j in np.ndindex(2,2)] 

другой

list(a.reshape(-1,3))