如何在numpy中的大张量上加快logsumexp？

Question

我在非常大的张量上在python中执行以下操作。

# p: 10000x1
# T,W: 400x400x10000
# Kp: 400x400
T = np.tile(p, (n, n, 1))
Kp = np.triu(scipy.special.logsumexp(T + W, axis=-1), k=1)

我知道我可以使用GPU并通过PyTorch的logsumexp大大提高速度，但是我现在想尝试优化CPU绑定版本。