在TensorFlow中收缩张量指数

Question

Suppose I have a tensor objects represented as $A_{i_0 i_1 ... i_k ... i_N}$ and $B_{j_0 j_1 ... j_p ... j_M}$ (in Tensorflow the would have shapes of N and M dimensions respectively). I want to create a contraction over dimensions k and p , so basically I want to create

$$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}....$$

What would be the right ops for this case?

Answer 1

Starting from Tensorflow 11 you can use einsum to do that.

So assuming A_ijkl and B_mnp and supposing you want to contract j with p you can do:

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