Numpy frompyfunc传入(2,4)与(2,1)数组时执行8次而非2次的问题

问题描述 #

我定义了一个接收两个Numpy数组并返回布尔值的函数x，使用np.frompyfunc将其包装为wrappedFunc。传入形状为(2,4)的arg_a和形状为(2,1)的arg_b后，函数被执行了8次（打印a和b各8次），但预期是执行2次：每次接收arg_a的一行数组和对应的arg_b元素。实际输出中a被打印为0-7，b的两个元素各重复打印4次；预期是第一次调用a为[0,1,2,3]、b为对应数组，第二次调用a为[4,5,6,7]、b为另一个数组。

代码示例 #

import numpy as np

def x(a,b):
    print(a)
    print(b)
    # Some computation...
    return boolean_value

wrappedFunc = np.frompyfunc(x,nin=2,nout=1)

arg_a = np.arange(8).reshape(2,4)
# arg_b is a numpy array having shape (2,1)
arg_b = np.array((np.array([[0, 1, 0],
                            [0, 0, 0],
                            [1, 0, 0],
                            [1, 1, 0]]),
                  np.array([[0., 1., 0.],
                            [0., 0., 0.],
                            [1., 0., 0.],
                            [1., 1., 0.],
                            [0.5, 0.5, 0.]])), dtype=object).reshape(2, 1)

问题原因 #

• np.frompyfunc的核心逻辑是逐元素广播执行：它会按照NumPy的广播规则，将输入数组展开到单个标量元素的级别进行配对调用，而非保留原数组的维度做批量处理。
• 你的arg_a是(2,4)的数值型数组，包含8个标量元素；arg_b是(2,1)的object型数组，每个元素是子数组。当frompyfunc处理时，arg_b会被广播为(2,4)的形状（每个子数组重复4次），最终和arg_a的8个标量一一配对，导致函数被调用8次。

解决方案 #

如果要实现按行批量调用的逻辑，不要使用np.frompyfunc，可以选择以下两种方式：

方式1：使用np.vectorize指定signature参数 #

np.vectorize支持通过signature参数定义输入输出的形状，从而实现按维度批量处理：

import numpy as np

def x(a,b):
    print(a)
    print(b)
    # 示例返回布尔值
    return True

# signature指定：第一个输入是一维数组(n)，第二个输入是单元素()，输出是标量()
vectorizedFunc = np.vectorize(x, signature='(n),()->()')

arg_a = np.arange(8).reshape(2,4)
arg_b = np.array((np.array([[0, 1, 0],
                            [0, 0, 0],
                            [1, 0, 0],
                            [1, 1, 0]]),
                  np.array([[0., 1., 0.],
                            [0., 0., 0.],
                            [1., 0., 0.],
                            [1., 1., 0.],
                            [0.5, 0.5, 0.]])), dtype=object).reshape(2, 1)

result = vectorizedFunc(arg_a, arg_b)

此代码会执行2次调用，每次传入arg_a的一行数组和对应的arg_b元素。

方式2：直接使用Python循环 #

逻辑更直观，适合简单的批量处理场景：

import numpy as np

def x(a,b):
    print(a)
    print(b)
    # 示例返回布尔值
    return True

arg_a = np.arange(8).reshape(2,4)
arg_b = np.array((np.array([[0, 1, 0],
                            [0, 0, 0],
                            [1, 0, 0],
                            [1, 1, 0]]),
                  np.array([[0., 1., 0.],
                            [0., 0., 0.],
                            [1., 0., 0.],
                            [1., 1., 0.],
                            [0.5, 0.5, 0.]])), dtype=object).reshape(2, 1)

result = []
# 遍历arg_a的行和arg_b的扁平化元素
for a_row, b_val in zip(arg_a, arg_b.flatten()):
    res = x(a_row, b_val)
    result.append(res)
# 转换为NumPy布尔数组
result = np.array(result, dtype=bool)

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