是否可以使用EinsumDense代替多个平行密集层？

Question

我有一个shape(None, 20, 250)的输入，其中20是我的上下文窗口，250是我的嵌入维度

我想应用不同的密度250->；上下文窗口中的每个元素对应250。 以下代码可以正常工作并执行我想要的操作，但不使用Einsum：

x = tf.unstack(x, axis=1)
x = [layers.Dense(250)(i) for i in x]
x = tf.stack(x, axis=1)

如果我正确理解了Einsum，那么以下代码也应该这样做：

einsum = EinsumDense('abc,cbd->abd', output_shape=[20, 250], bias_axes='bd')
x = einsum(x)

EinsumDense的文档说明ab,bc->ac相当于一个密集层，所以我认为这应该是可行的。 然而，当我这样做的时候，我得到了截然不同的结果

我做错了什么