在固定维度的较小块中拆分3D数组的Pythonic方法

我想在不使用循环的情况下实现以下例程，只需使用 Numpy 函数或 ndarray 方法。这是代码：

def split_array_into_blocks( the_array, block_dim, total_blocks_per_row ):

    n_grid = the_array.shape[0]
    res    = np.empty( (n_grid, total_blocks_per_row, total_blocks_per_row, block_dim, 
    block_dim ) )

    for i in range( total_blocks_per_row  ):
        for j in range( total_blocks_per_row ):
        
            subblock          = the_array[ :, block_dim*i:block_dim*(i+1), block_dim*j:block_dim*(j+1) ]   
            res[ :, i,j,:,: ] = subblock 

    return res

我尝试过“重塑”方法，以便：

the_array = the_array.reshape( ( n_grid, total_blocks_per_row, total_blocks_per_row, block_dim, block_dim) )

但这似乎以某种方式改变了元素的顺序，并且块需要像在例程中一样存储。任何人都可以提供一种方法来做到这一点，并简要解释为什么 reshape 方法在这里给出不同的结果？（也许我还缺少使用 np.transpose() ？）

编辑：我想出了这个替代实现，但我仍然不确定这是否是最有效的方法（也许有人可以在这里说明）：

def split_array_in_blocks( the_array, block_dim, total_blocks_per_row ):
       
    indx   = [ block_dim*j for j in range( 1, total_blocks_per_row ) ] 

    the_array = np.array( [ np.split(  np.split( the_array, indx, axis=1 )[j], indx, axis=-1 ) for j in range( total_blocks_per_row ) ] )
    the_array = np.transpose( the_array, axes=( 2,0,1,3,4 ) )


    return the_array

示例：这是两个实现的最小工作示例。我们想要的是，从维度 Nx3 MX3 M的初始“立方体”分解为块 NxMxMx3x3，它们是原始块的分块版本。通过上述两种实现，可以检查它们是否给出相同的结果；问题是如何以有效的方式（即，无循环）实现这一目标

import numpy as np


def split_array_in_blocks_2( the_array, block_dim, total_blocks_per_row ):
    
    n_grid = the_array.shape[0]
    res    = np.zeros( (n_grid, total_blocks_per_row, total_blocks_per_row, block_dim, block_dim ), dtype=the_array.dtype )
    
    for i in range( total_blocks_per_row ):
        for j in range( total_blocks_per_row ):            
            subblock          = the_array[ :, block_dim*i:block_dim*(i+1), block_dim*j:block_dim*(j+1) ]   
            res[ :, i,j,:,: ] = subblock 
    
    return res


def split_array_in_blocks( the_array, block_dim, total_blocks_per_row ):
           
    indx   = [ block_dim*j for j in range( 1, total_blocks_per_row ) ] 
    
    the_array = np.array( [ np.split(  np.split( the_array, indx, axis=1 )[j], indx, axis=-1 ) for j in range( total_blocks_per_row ) ] )
    the_array = np.transpose( the_array, axes=( 2,0,1,3,4 ) )
    
    
    return the_array

A = np.random.rand( 1001, 63, 63 )
n = 3
D = 21
from time import time

ts = time()
An = split_array_in_blocks( A, n, D )

t2 = time()

Bn = split_array_in_blocks_2( A, n, D )

t3 = time()
print( t2-ts )
print(t3-t2)
print(np.allclose( An, Bn ))