TensorFlow中的收縮張量索引

Question

假設我有一個張量對象，表示爲$A_{i_0 i_1 ... i_k ... i_N}$和$B_{j_0 j_1 ... j_p ... j_M}$（在Tensorflow中將分別具有N和M維的形狀）。我想創造超過尺寸ķ和p收縮，所以基本上我想創建

$$A_{i_0 i_1 ... 0 ... i_N} B_{j_0 j_1 ... 0 ... j_M} + A_{i_0 i_1 ... 1 ... i_N} B_{j_0 j_1 ... 1 ... j_M} + A_{i_0 i_1 ... 2 ... i_N} B_{j_0 j_1 ... 2 ... j_M}....$$ 

什麼會是這種情況下正確的OPS？

Answer 1

從Tensorflow 11開始，您可以使用einsum來做到這一點。

因此，假設A_ijkl和B_mnp以及假如你希望合同j與p你可以這樣做：

import tensorflow as tf 
tf.einsum("ijkl,mnj->iklmn", A,B)