Numpy 3D матричное умножение

Question

У меня есть 2 матрицы A (форма 10x10x36) и B (форма 10x27x36). Я хотел бы умножить последние 2 оси и суммировать результат по оси 0 так, чтобы результат C имел форму 10x27. Вот в настоящее время, как я это делаю

C = []
for i in range(A.shape[0]):
    C.append(np.matmul(A[i], B[i].T))
C = np.sum(np.array(C), axis=0)

Я хочу добиться этого в векторизации, но не могу понять, как это сделать. Я проверил np.einsum , но пока не уверен, как применить его для достижения результата. Любая помощь будет оценена. Спасибо!

Answer 1

Здесь тот же результат, используя np.einsum:

r1 = np.einsum('ijk,ilk->jl', A, B)

Однако на моей машине реализация для l oop работает почти в 2 раза быстрее:

def f(A,B):
    C = []
    for i in range(A.shape[0]):
        C.append(np.matmul(A[i], B[i].T))
    return np.sum(np.array(C), axis=0)

%timeit np.einsum('ijk,ilk->jl',A,B)
102 µs ± 3.79 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit f(A,B)
57.6 µs ± 1.7 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Answer 2

matmul поддерживает укладку. Вы можете просто сделать:

 (A@B.transpose(0,2,1)).sum(0)

Проверки (C генерируется с использованием OP l oop):

np.allclose((A@B.transpose(0,2,1)).sum(0),C)
# True
timeit(lambda:(A@B.transpose(0,2,1)).sum(0),number=1000)
# 0.03199950899579562
# twice as fast as original loop

Answer 3

Вы также можете попробовать следующее, используя понимание списка. Это немного более кратко, чем то, что вы сейчас используете.

C=np.array([A[i] @ B.T[:,:,i] for i in range(10)]).sum(0)