First vectorise foo()
, i.e. modify foo()
in a way that it can correctly operate on an array of shape (N, A, B)
, returning an array of shape (A, B)
. This step is usually the difficult one. How this is done entirely depends on what foo()
does. For the given example, it's very easy to do:
def foo(arr):
return np.sum(arr, axis=0)
Now, use broadcasting rules to create a (N, A, B)
array containing all the vector differences, and pass it to foo()
:
foo(a[:, :, np.newaxis] - b[:, np.newaxis])
与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…