我想对函数f(a, b)
进行向量化,以便当我输入a和b作为两个向量时,返回组合的张量。这是一个说明性的例子:
import numpy as np
def tester(a, b):
mysumm = 0.
for ii in range(a):
for jj in range(b):
mysumm += a * b
return mysumm
tester = np.vectorize(tester)
x, y = [2, 4], [3, 5, 8]
print(tester(x, 3)) # [ 36. 144.]
print(tester(x, 5)) # [100. 400.]
print(tester(x, 8)) # [ 256. 1024.]
print(tester(2, y)) # [ 36. 100. 256.]
print(tester(4, y)) # [ 144. 400. 1024.]
print(tester(x, y)) # ValueError: operands could not be broadcast together with shapes (2,) (3,)
我希望tester(x, y)
调用返回2x3矩阵,类似于[[ 36. 100. 256.], [ 144. 400. 1024.]]
,我很惊讶这不是默认行为。
如何使vecotirzed函数返回输入向量的可能组合的张量?
答案 0 :(得分:2)
您可以使用np.ix_
:
>>> import functools
>>>
>>> def tensorize(f):
... fv = np.vectorize(f)
... @functools.wraps(f)
... def ft(*args):
... return fv(*np.ix_(*map(np.ravel, args)))
... return ft
...
>>> tester = tensorize(tester)
>>> tester(np.arange(3), np.arange(2))
array([[0., 0.],
[0., 1.],
[0., 4.]])