python - Python：同时对两个多维数组进行矢量化函数

Question

在 python 中，对单个数组f(x)上的标量函数进行矢量化很简单：只需使用. 但是假设我有两个（或原则上是多个）具有相同形状的数组，并且我想对一个带有标量的函数同时在两者上进行矢量化。类似的东西，它将再次返回一个具有共同维度的对象。xa1f(a1)a1, a2Nx3g(x,y)x,yg(a1,a2)N

编辑：如果a1,a2都是一维的，这变得微不足道。我们使用如下所述的简单广播。但是，对于多维数组，答案对我来说并不明显。那么，我该怎么做，最好使用numpy？

示例（已编辑）：

a1 = np.array of size 20x3 # so that each row is a 3-vector
a2 = np.array of size 20x3 # ditto

def f(x, y): # acts on each element 
   ... complicated function, using other global variables ...
   return ... (scalar)

如果没有矢量化，我需要f单独循环所有 20 行，并获得输出长度为 20 的矢量：

result = []
for i, elem in a1:
    result.append(f(elem, a2[i]))
result = np.array(result)

但是，我想消除for循环，并使用 numpy 矢量化有一个语句。原因是能够使用jax( https://github.com/google/jax ) 的 numpy 包装器，然后在 GPU 上加速它。像天真的东西

result = f(a1, a2) 

不起作用。那么正确的语法是什么？

标签： pythonnumpyvectorization

Answer 1

它可能取决于您需要执行的操作，如果它是一个简单的总和，那么以下将起作用：

import numpy as np
a = np.arange(3*2*20).reshape((20,3,2))
b = np.arange(2*20).reshape((20,2))

res = (a.transpose((1,2,0))+b.transpose((1,0))).transpose((2,0,1))
print(a[0],b[0])

[[0 1]
 [2 3]
 [4 5]] [0 1]

print(res[0])

[[0 2]
 [2 4]
 [4 6]]

首先，输入数据被转置，以便在广播操作中涉及正确的维度。求和后，输出转回。