1
$\begingroup$

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?

$\endgroup$

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.