使用numpy和functools向量化偏函数时的异常行为排查

问题背景 #

尝试向量化一个接收两个列表参数的偏函数，该函数通过zip处理成对元素，但出现异常行为。

代码示例 #

import functools
import numpy as np

def f(l1,l2):
    l1 = l1 if isinstance(l1,list) or isinstance(l1,np.ndarray) else [l1]
    l2 = l2 if isinstance(l2,list) or isinstance(l2,np.ndarray) else [l2]
    for e1,e2 in zip(l1,l2):
        print(e1,e2)

f(['a','b'],[1,2])

fp = functools.partial(f,l1=['a','b'])
fp(l2=[1,2])

fv = np.vectorize(fp)
fv(l2=np.array([1,2]))

Jupyter Notebook输出 #

a 1
b 2

a 1
b 2

a 1
a 1
a 2
array([None, None], dtype=object)

补充尝试代码 #

fpv(l2=[np.array([1,2]), np.array([3,4])])

补充输出 #

a 1
a 1
a 2
a 3
a 4

核心问题 #

1. 函数f开头的类型检查是必需的，因为np.vectorize会自动完全扁平化输入（否则会抛出int32 not iterable异常），是否有方法可以避免这种类型检查？
2. 偏函数fp被向量化后，输出不符合预期，不清楚NumPy的执行逻辑，包括最终的空数组输出。无论如何嵌套[1,2]，输出始终一致，如何修改代码使向量化函数fv的行为与fp一致？

问题解答 #

问题1：避免类型检查的方法 #

np.vectorize默认会遍历输入数组的每个标量元素调用目标函数，导致输入被完全扁平化，单个数值传入后触发迭代报错。要跳过类型检查，只需通过signature参数告诉np.vectorize：把整个数组/列表作为单个参数传入，而非拆分元素。

修改后代码：

# 直接移除类型检查逻辑
def f(l1,l2):
    for e1,e2 in zip(l1,l2):
        print(e1,e2)

fp = functools.partial(f,l1=['a','b'])
# signature定义输入为单个一维序列，输出无返回值
fv = np.vectorize(fp, signature="()->()")
fv(l2=np.array([1,2]))

此时np.vectorize会将np.array([1,2])整体传给l2，无需额外类型转换。

问题2：让向量化函数与偏函数行为一致 #

当前异常输出的原因是：np.vectorize默认拆分输入数组的每个标量（1和2）分别传入fp，每次传入单个数值时，f的类型检查会将其转为单元素列表，和l1=['a','b']``zip后只能输出一组a x；同时f无返回值，np.vectorize收集到None后生成对应形状的数组。

要让fv和fp行为一致，核心还是用signature参数约束输入传递方式：

• 若处理单个列表/数组：用signature="()->()"，确保整体传入
• 若批量处理多个列表/数组（如补充尝试的场景）：用signature="(n)->()"，表示输入是长度为n的一维序列，每次传递完整序列

批量处理示例代码：

fv = np.vectorize(fp, signature="(n)->()")
fv(l2=[np.array([1,2]), np.array([3,4])])

输出会和连续调用fp(l2=[1,2])、fp(l2=[3,4])完全一致：

a 1
b 2
a 3
b 4

额外说明 #

• np.vectorize本质是Python循环的封装，并非真正的向量化运算，性能和纯Python循环接近，不要依赖它提升计算速度。
• 若场景是批量处理成对序列，直接用Python循环遍历输入列表调用fp，可能比np.vectorize更直观高效。

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