# Is it possible to use EinsumDense instead of multiple parallel Dense layers?

I have an input of shape (None, 20, 250), where 20 is my context window and 250 my embedding dimension.

I want to apply a different dense 250 -> 250 for each element in the context window. The following code works and does what I want, but does not use Einsum:

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


If I understand Einsum correctly, the following code should do the exact same:

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


The documentation for EinsumDense states that ab,bc->ac would be equivalent to a dense layer, so I figured this should work. However I'm getting drastically different results when doing this.

What am I doing wrong here?