如何将某个函数应用于沿两个张量的维度的所有组合

Question

我正试图在torch中实现这样的目标。

我有两个形状为[X, 4]和[Y, 4]的张量。这个4s基本上是某事物的4个坐标。因此，对于X和Y的每个组合(两个长度为4的向量(，我想应用一些函数(例如元素平均(，并形成形状[X, Y, 4]的结果。

如何做到这一点？

所谓元素平均，我指的是这个操作，

[2 4 6 8] OP [8 6 4 2] = [5 5 5 5]

但它可以是任意操作。

注意：我能够使用循环来解决它，但正在寻找一个矢量化的解决方案。

Answer 1

您的答案取决于函数。例如，对于元素平均值：

np.mean([x[:,None,:],y[None,...]], axis=0)

或对于einsum:

np.einsum('ij,kj->ikj',x,y)

或求和：

x[:,None,:]+y[None,...]

如果您正确地实现了CUSTOM功能(elementwise(，您可以使用广播来完成这项工作：

np.frompyfunc(myfunc,2,1)(x[:,None,:],y[None,...]) 

示例：

x = array([[0, 1, 2, 3],
[4, 5, 6, 7]])
y = array([[100, 101, 102, 103],
[104, 105, 106, 107],
[108, 109, 110, 111]])
#np.mean([x[:,None,:],y[None,...]], axis=0)
array([[[50, 51, 52, 53],
[52, 53, 54, 55],
[54, 55, 56, 57]],
[[52, 53, 54, 55],
[54, 55, 56, 57],
[56, 57, 58, 59]]])
#np.einsum('ij,kj->ikj',x,y)
array([[[  0, 101, 204, 309],
[  0, 105, 212, 321],
[  0, 109, 220, 333]],
[[400, 505, 612, 721],
[416, 525, 636, 749],
[432, 545, 660, 777]]])

#x[:,None,:]+y[None,...]
array([[[100, 102, 104, 106],
[104, 106, 108, 110],
[108, 110, 112, 114]],
[[104, 106, 108, 110],
[108, 110, 112, 114],
[112, 114, 116, 118]]])
def myfunc(x, y):
return x+y
#np.frompyfunc(myfunc,2,1)(x[:,None,:],y[None,...]) 
array([[[100, 102, 104, 106],
[104, 106, 108, 110],
[108, 110, 112, 114]],
[[104, 106, 108, 110],
[108, 110, 112, 114],
[112, 114, 116, 118]]], dtype=object)