python - Efficiently re-stacking a numpy ndarray -

i have problem re-stacking numpy ndarrays in manner. want re-stack inner 2d matrices columns adjacent each other. example, if input:

>>> k.shape (2, 3, 3, 2) >>> k array([[[[ 7.,  6.],          [ 7.,  5.],          [ 5.,  5.]],          [[ 2.,  7.],          [ 5.,  2.],          [ 7.,  1.]],          [[ 9.,  1.],          [ 7.,  1.],          [ 2.,  6.]]],          [[[ 5.,  8.],          [ 9.,  6.],          [ 3.,  7.]],          [[ 8.,  2.],          [ 2.,  8.],          [ 4.,  4.]],          [[ 8.,  9.],          [ 5.,  9.],          [ 2.,  4.]]]], dtype=float32) 

i know need know np.hstack, don't know how well. here's how i'm doing it, , want like:

>>> np.array([np.hstack(j) j in k]) array([[[ 7.,  6.,  2.,  7.,  9.,  1.],     [ 7.,  5.,  5.,  2.,  7.,  1.],     [ 5.,  5.,  7.,  1.,  2.,  6.]],     [[ 5.,  8.,  8.,  2.,  8.,  9.],     [ 9.,  6.,  2.,  8.,  5.,  9.],     [ 3.,  7.,  4.,  4.,  2.,  4.]]], dtype=float32) 

you can imagine this: have 2 images of dimension 3x2 , each image has 3 channels. want horizontally stack channels of each image next each other.

i apologise if example little weird, data working not simpler this.

list comprehension in python slow. there faster way this?

thanks!

we can permute axes np.transpose , reshape -

arr.transpose(0,2,1,3).reshape(2,3,-1)