numpy函数的力输出形状

Question

我有一个函数，它将2个numpy数组作为输入并返回一个数字（实函数比下面的示例复杂得多）：

def the_func(x,y): 
    return np.sum(x)*np.sum(y)

>>> the_func([2,3,4],[10,11,12])
297

到目前为止，一切顺利

但是

>>> the_func([[1,2,3],[4,5,6]],[[1,2,3],[4,5,6]])
441

而不是预期的[36,225]

因此我尝试vectorize初始函数：

vfunc = np.vectorize(the_func)

但是

>>> vfunc([[1,2,3],[4,5,6]],[[1,2,3],[4,5,6]])
array([[ 1,  4,  9],
       [16, 25, 36]])

如何让该函数以我希望看到的形状返回结果（在上一个示例中[36,225]=[(1+2+3)*(1+2+3),(4+5+6)*(4+5+6)]）

Answer 1

您应该使用axis关键字参数来指定它应该和的维度。在您的示例中，将函数更改为：

def the_func(x,y): 
    return np.sum(x, axis=1)*np.sum(y, axis=1)

在示例输入上运行它会得到结果[36 255]