Gibberish with tensorflow 2.0.0b0

[xsys-technology]

Thank you for the code and for wading through all of those einsum's to get things working!

I'm trying to generate some text using tensorflow 2.0 and have thus far only been able to produce gibberish.

In order to get the code running with the 2.0 api I had to make a few minor modifications: 1) change tf -> tf.compat.v1 in each place the compiler complained, and 2) use tf.keras.layers.LayerNormalization instead of tf.contrib.layers.layer_norm in parts of the attention and ffn code: (e.g.

#output = tf.contrib.layers.layer_norm(output + inp, begin_norm_axis=-1, scope='LayerNorm')
ln = tf.keras.layers.LayerNormalization()
with tf.compat.v1.variable_scope("LayerNorm"):
    output = ln(output + inp)

The documentation for LayerNormalization implies the default behavior is equivalent to using begin_norm_axis=-1 so I don't think this is the issue.

Any ideas to remedy the gibberish?