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?