如何在numpy中沿指定轴应用接收两个数组参数的自定义函数

方法1：拆分维度+原生向量化改造（性能最优） #

你可以先直接拆分axis=2维度上的两个数组，利用numpy的广播特性直接批量处理，不需要额外调用遍历类接口：

import numpy as np

# 生成示例数组
x = np.random.rand(3, 10, 2, 6)

# 拆分axis=2的两个数组，形状均为 (3, 10, 6)
arr1, arr2 = x[..., 0, :], x[..., 1, :]

# 改造你的函数f，所有操作使用numpy原生向量化算子，天然支持前导批量维度
# 示例1：输出两个数组对应元素和，输出形状为 (3,10,6)
def f(arr1, arr2):
    return arr1 + arr2

# 示例2：输出两个数组各自求和后的和，输出形状为 (3,10)
def f(arr1, arr2):
    return np.sum(arr1, axis=-1) + np.sum(arr2, axis=-1)

result = f(arr1, arr2)
# 若需要输出为 [3,10, dim] 格式，dim为1时可扩展维度：result = result[..., np.newaxis]

方法2：多输入向量包装（适配不可修改的单组输入函数） #

如果你的函数f只能接收1D数组输入、无法修改为支持批量维度，可以用np.vectorize指定输入输出签名实现适配：

# 你的原始仅支持1D输入的函数
def f(arr1, arr2): 
    # arr1 形状为 (6, )，arr2 形状为 (6, )
    return np.sum(arr1) + np.sum(arr2)

# 包装为支持批量输入的向量函数，签名定义为：两个长度为n的1D输入，输出一个标量
f_vec = np.vectorize(f, signature='(n),(n)->()')

result = f_vec(arr1, arr2)
# 输出形状为 (3,10)，需要扩展最后一维可添加 result = result[..., np.newaxis]

你之前了解的numpy.apply_along_axis仅支持单1D切片作为输入，确实不适用当前需要传入两个切片参数的场景，上述两种方案均可满足需求，优先使用方法1性能更高。

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