如何设置NumPy vectorize的signature参数以实现指定输出？

先明确我的函数定义、尝试的向量化写法、输入数组及期望输出：

函数与输入输出 #

import numpy as np

# 定义的原始函数
def fun1(a):
    b = a[0] + a[1]
    return(b)

# 尝试的向量化写法
fun2 = np.vectorize(fun1, signature='(n,m)->(n,1)')

# 输入数组
input_arr = np.array([
[1 ,2],
[3, 4],
[5, 6]])

# 期望输出
# np.array([[3],[7],[11]])

我知道用np.sum(input_arr, axis=1).reshape(-1,1)能得到结果，但重点是搞懂signature的正确用法，这样才能处理更复杂的函数。

为什么原signature不对 #

np.vectorize的signature参数，描述的是单个函数调用的输入输出形状映射，而非整个输入/输出数组的形状。

原写法'(n,m)->(n,1)'的问题在于：它要求把整个(n,m)的二维数组作为单个参数传给fun1，但fun1的逻辑是取a[0]+a[1]——如果a是二维数组，a[0]是第一行、a[1]是第二行，相加会得到[4,6]，完全不是我们要的逐行计算结果。

正确的signature设置 #

我们的需求是：对输入数组的每一行（形状为(m,)的一维数组）调用fun1，返回一个标量，最后把所有标量拼成列向量。

方式一：输出标量后reshape #

设置signature='(m)->()'，表示单个输入是(m,)的一维数组，输出是标量()：

fun2 = np.vectorize(fun1, signature='(m)->()')
result = fun2(input_arr).reshape(-1, 1)
print(result)
# 输出：
# [[ 3]
#  [ 7]
#  [11]]

方式二：直接输出单元素数组 #

如果想让vectorize直接返回二维列向量，可以把signature设为'(m)->(1,)'，表示单个输入返回(1,)的一维数组，最终拼接成(n,1)的二维数组：

fun2 = np.vectorize(fun1, signature='(m)->(1,)')
result = fun2(input_arr)
print(result)
# 输出：
# [[ 3]
#  [ 7]
#  [11]]

核心总结 #

设置signature的关键步骤：

1. 明确原始函数fun1能处理的单个输入的形状（比如这里是长度为2的一维数组(m,)）
2. 明确原始函数返回的单个输出的形状（标量()或单元素数组(1,)）
3. 用'输入形状->输出形状'的格式写出signature，np.vectorize会自动遍历输入中对应形状的子块，批量调用函数并拼接结果

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