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?


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.