Question

假设我们有一个形状为 (N, K) 的数组 X 和一个形状为 (N, P, K) 的数组 Y。

目标是为沿轴 0 的每个 i=1,...,N 有效地计算一些函数 f(x_i, y_i)。请注意，x_i 的形状为 (1, K)，y_i 的形状为 (P, K ).

我的简单实现只是第一个轴上的 for 循环。似乎应该有某种方法以更矢量化的方式来完成它，或者使用广播规则一步一步地在完整的数组 X、Y 上计算 f - 而不是对每个 i 逐一进行。关于实现此逻辑的更有效方法的任何建议？

import numpy as np

N = 10
P = 7
K = 5

# Shape (N, K) array
X = [np.arange(0,K) for i in range(N)]
X = np.stack(X)

# Shape (N, P, K) array
Y = [(-1)**k for k in range(N*P*K)]
Y = np.array(Y).reshape(N,P,K)

# Some function f: 
# For simplicity, just adds x+y with shape broadcasting between a (1,K) array and (*, K) array
def f(x_i, y_batch):
    l = x_i + y_batch
    return l

# Goal: Efficiently compute some function f of the ith element (axis 0) of X and Y for each i=1,...,N, 

# Example: Naive Implementation
results = []
for i in range(N):
    x = X[i,:]
    y = Y[i,:,:]
    results.append( f(x,y) )

在不同秩数组之间高效切片和应用函数

0 个答案: