如何向函数传入多NumPy数组并返回多输出？行向量求和优化

你遇到的问题是因为直接取x[0] + x[1]是对行进行元素级相加，而非对每行内部的元素求和。以下是几种无需循环的高效实现方式：

方法1：使用numpy.sum()指定轴参数 #

这是最简洁的方式，通过axis=1指定按行求和：

import numpy as np
inps = np.array([[1, 2], [3, 4], [5, 6]])
sum_func = lambda x: x.sum(axis=1)
print(sum_func(inps))
# >>> [ 3  7 11]

方法2：直接索引每行元素相加 #

如果你的数组固定是2列，也可以直接取每行的两个元素做元素级加法：

import numpy as np
inps = np.array([[1, 2], [3, 4], [5, 6]])
sum_func = lambda x: x[:, 0] + x[:, 1]
print(sum_func(inps))
# >>> [ 3  7 11]

方法3：使用numpy.einsum() #

einsum可以灵活指定求和维度，'ij->i'表示对二维数组的第二个维度（列）求和，保留第一个维度（行）：

import numpy as np
inps = np.array([[1, 2], [3, 4], [5, 6]])
sum_func = lambda x: np.einsum('ij->i', x)
print(sum_func(inps))
# >>> [ 3  7 11]

原代码问题解释 #

你原来的lambda x : x[0] + x[1]中，x[0]是数组的第一行[1,2]，x[1]是第二行[3,4]，两者相加是元素级别的对应相加，得到[1+3, 2+4] = [4,6]（你描述的[3,3]可能是笔误），这和你期望的每行内部求和逻辑不符。

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