Transformer example masking is incorrect when using Keras MultiHeadedAttention

Referring to the Transformer example.

The script probably works correctly as is. However, there is one part where you mention if you want to use tf.keras MultiHeadedAttention layer, instead of the custom one created here, the look-ahead mask has to be inverted (because the layer attends to 1's in the padding mask, not 0's). What is missed is the fact that the padding mask also has to be inverted, since it is used for attention, e.g. in the DecoderLayer. That change would also necessitate some changes to the loss function and* Transformer.create_masks() (and anywhere else the padding mask is used). Having two versions of all these functions might be too big of a hassle, but at least commenting where each change should be made would be useful. Even better would be to update the custom multi headed attention so that the padding is consistent with Keras's approach.

* I realize now the loss mask is computed directly from the target, and not passed to the function, so that part is still fine

Start with the Transformer example notebook, focusing on DecoderLayer, Transformer.create_masks(), and the custom and tf.keras MultiHeadedAttention implementations. Check how look-ahead and padding masks are represented and used, then verify the example works consistently with the Keras layer and that the affected mask-related behavior is documented or corrected.