Better handling of padded values

[ivyleavedtoadflax]

Padded values are currently handled in a custom way, but could probably be much better handled using tensorflow masking, see: https://www.tensorflow.org/api_docs/python/tf/keras/layers/Embedding, and
https://www.tensorflow.org/guide/keras/masking_and_padding