3D numpy矩阵最后一维作函数入参的多进程计算实现方法

我有一个存储序列数据的DataFrame，需要对其中的序列做两两比对计算，待处理序列存储在x_test Series中，结构如下：

x_test
0    [0.9948399033268663, 0.9948399033268663, 0.995...
1    [1.0005639107922013, 1.00041937494038, 1.00020...
2    [0.9975316557982873, 0.9975316557982873, 0.997...
Name: cycle_info_scaled, dtype: object

距离比对使用vec_compute_distance函数实现，该函数接收2个列表作为入参返回计算得到的距离值，代码如下：

def vec_compute_distance(*args):
    args = args[0]
    if len(args) != 2:
        print(f"args lens is {len(args)}, Should take 2 values")
        return 0
    else:
        return _ucrdtw.ucrdtw(args[0], args[1], 0.05, False)[1]

为实现两两配对映射计算，我基于该Series构造3D矩阵：

grid_test = np.asarray(np.meshgrid(x_test, x_test)).T

其中grid_test的维度为：

grid_test.shape
(3, 3, 2)

我的目标是将该3x3矩阵每个位置的最后一维（共2个元素）作为函数入参完成映射计算。由于待处理数据量极大，我需要将原有双层for循环的实现改为适配map的写法，以支持multiprocessing多进程并行计算，但改写过程中触发如下报错：

ValueError                                Traceback (most recent call last)
Input In [26], in <cell line: 3>()
      2 grid_test = np.asarray(np.meshgrid(x_test, x_test)).T
      3 for i in range(grid_test.shape[0]):
----> 4     grid_test[i] = list(map(vec_compute_distance, grid_test[i]))

ValueError: could not broadcast input array from shape (3) into shape (3,2)

可正常运行但效率过低的双层循环实现如下：

for i in range(grid_test.shape[0]):
    for j in range(grid_test.shape[1]):
        grid_test[i,j] = vec_compute_distance(grid_test[i,j])

寻求该问题的可行解决方案，要求实现可适配map接口、支持多进程并行，满足大规模数据的高效处理需求。

内容的提问来源于stack exchange，提问作者josseossa