Блог

Предположим, у нас есть массив X с формой (N, K) и массив Y с формой (N, P, K).

Цель состоит в том, чтобы эффективно вычислить некоторую функцию f(x_i, y_i) для каждого i=1,…, N вдоль оси 0. Обратите внимание, что x_i имеет форму (1, K), а y_i имеет форму (P, K).

Наивная реализация, которую я имею, — это просто цикл for по первой оси. Похоже, что должен быть какой — то способ сделать это более векторизованным способом или использовать правила широковещания для вычисления f на полных массивах X, Y за один шаг-вместо того, чтобы делать это для каждого 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) )