Тензорная операция в Tensorflow

Question

У меня есть 2 тензора, тензор A и тензор B, как показано на рисунке. Я хочу взять каждый вектор 1x16 из тензора B и выполнить точечное произведение с соответствующей партией (ось 0 - размер партии в каждом тензоре) тензора A вдоль осей 1 и 2 из A. Как это сделать в тензорном потоке. Любая помощь с кодом будет принята с благодарностью. Спасибо

Я пробовал einsum, но я в замешательстве.

Ожидаемая форма вывода должна быть (3,2,28,28).

Answer 1

Вы можете сделать это так:

import tensorflow as tf

A = tf.placeholder(tf.float32, [3, 28, 28, 16])
B = tf.placeholder(tf.float32, [3, 2, 16])
A_exp = A[:, tf.newaxis]
B_exp = B[:, :, tf.newaxis, tf.newaxis]
C = tf.reduce_sum(A_exp * B_exp, axis=-1)
print(C)
# Tensor("Sum:0", shape=(3, 2, 28, 28), dtype=float32)