TensorFlow1.10中如何用矩阵乘法实现(Batch,a,b)乘(Batch,b)输出(Batch,a)张量

方案1：使用tf.matmul（维度适配后批量相乘） #

该方案通过调整张量形状适配批量矩阵乘法的规则，实现需求：

• 给形状为(Batch, b)的t2扩展最后一个维度，得到形状为(Batch, b, 1)的张量，满足矩阵乘法的维度要求
• 调用tf.matmul执行批量矩阵乘法，得到形状为(Batch, a, 1)的中间结果
• 调用tf.squeeze去掉最后一个长度为1的维度，得到最终(Batch, a)的输出

示例代码：

import tensorflow as tf
# TF1.10版本如需开启eager执行可添加下行
# tf.enable_eager_execution()

# t1、t2为对应形状的已定义张量
t2_expand = tf.expand_dims(t2, axis=-1)
res = tf.squeeze(tf.matmul(t1, t2_expand), axis=-1)

方案2：使用tf.einsum（更简洁，无需手动调整维度） #

einsum可以通过下标规则直接指定维度运算逻辑，不需要手动调整张量形状，代码可读性更高。
运算规则ijk,ik->ij的含义为：第一个张量有i、j、k三个维度（对应Batch、a、b），第二个张量有i、k两个维度（对应Batch、b），按k维度求和后输出i、j维度的结果，正好匹配需求。

示例代码：

import tensorflow as tf
# TF1.10版本如需开启eager执行可添加下行
# tf.enable_eager_execution()

res = tf.einsum('ijk,ik->ij', t1, t2)

以上两种方案在TensorFlow1.10版本中均可正常运行，代入你给出的示例输入可以得到目标输出结果。

内容的提问来源于stack exchange，提问作者Bowen Peng