在TensorFlow中收缩张量指数

Question

假设我有一个表示为$A_{i_0 i_1 ... i_k ... i_N}$和$B_{j_0 j_1 ... j_p ... j_M}$的张量对象（在Tensorflow中，它们分别具有N和M维的形状）。我想在维度 k 和 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}....$$

这种情况的正确操作是什么？

Answer 1

从Tensorflow 11开始，您可以使用einsum来执行此操作。

假设A_ijkl和B_mnp并假设您希望与j签订p合同，则可以执行以下操作：

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