python - pytorch中的张量幂和乘法

Question

我有一个矩阵 A 和一个大小为 b 的张量(1,3)- 所以一个大小为 3 的向量。

我要计算

C = b1 * A + b2 * A^2 + b3 * A^3其中^n是 A 的 n 次方。

最后，C 应该与 A 具有相同的形状。我怎样才能有效地做到这一点？

标签： pythonpytorchtensortorch

Answer 1

咱们试试吧：

A = torch.ones(1,2,3)
b_vals = torch.tensor([2,3,4])
powers = torch.tensor([1,2,3])

C = (A[...,None]**powers + b_vals).sum(-1)

输出：

tensor([[[12., 12., 12.],
         [12., 12., 12.]]])