numpy 中轴上的矩阵乘法

问题描述 投票:0回答:1

假设我有一个形状为 (B, N, N, 3, 3) 的数组

X
。我想向量化操作

X[:,0,...] @ X[:,1,...] ... @ X[:,N-1,...]

如何在 numpy 中对其进行矢量化?我不想使用 for 循环。

我尝试使用 einsum,但它对我不起作用,它最终对矩阵求和。

numpy pytorch matrix-multiplication jax numpy-einsum
1个回答
0
投票

我认为最简洁的表达方式是这样的:

functools.reduce(np.matmul, np.rollaxis(x, 1))

我怀疑这也和您希望使用 numpy 中可用的操作一样高效。

© www.soinside.com 2019 - 2024. All rights reserved.