New TF embeddings (cleaner and faster)
Author: jplu
Pull request: yes

# What does this PR do?

This PR propose a better implementation of the embedding layer for the BERT-Like TF models. Another benefit of this cleaning is a better computational performance:
```
model = TFBertForMaskedLM.from_pretrained("bert-base-cased")
cProfile.run("model(model.dummy_inputs)")

# current master
56150 function calls (55318 primitive calls) in 0.096 seconds

# with new embeddings implem
55732 function calls (54891 primitive calls) in 0.080 seconds
```

This new implementation should be compatible with the incoming rework of the resizing proposed in #9193. A similar work will be applied to `TFSharedEmbeddings` in a next PR.

All slow/quick tests passes.

EDIT: I don't know why Github has some issues to pin the reviewers, so pinging @LysandreJik @sgugger and @patrickvonplaten 

Comments:
I like this PR in general! 

Just wondering about two things:

1) Do we need this `get_config` function?
2) Not a huge fan of the `Add()` keras layer...does this really improve performance much?

Good point @LysandreJik! Basically here most of the models share the similar embedding computation that stay inside their respective file. What has been exported is just the specific computation, which means that `WordEmbeddings`, `PositionalEmbeddings` and `TokenTypeEmbeddings` are always the same doesn't matter who is using it.

The same logic that is currently applied to `TFSharedEmbeddings`.

> Just reviewed the general approach on one model for now and I have some questions before going further. If I understand correctly, the computation of the three different types of embeddings is split in three different ways to maximize the speedup but I wonder if it's documented from TF or just some tests on one particular setup. Before adding the extra complexity, I would like to be sure it brings a speedup on almost all possible environments (CPU, GPU, multi-GPU, TPU) without any loss in memory footprint (one-hot encoding the token type ids seems harmless, but we never know).

I basically took example on the official implementation of Transformer encoder available in the Google Repo https://github.com/tensorflow/models/tree/master/official/nlp/keras_nlp . After having done several experiments (only on CPU and GPU though), I end up to extract from this an optimal version for each embedding.

> As for putting those in modeling utils versus the model file, I agree with Lysandre that this breaks our philosophy of putting everything in each model file. I emitted the same reserves for TFSharedEmbeddings when it was introduced.

I don't mind to copy/paste the same layers in all the concerned files if it is the recommended way. @sgugger @LysandreJik Will you be more confident if I create a version for each model and add the comment `# copied from ....` everytime it is a strong copy/paste?

> I don't understand how it can be used above (line 420) in a tf.matmul if it's a layer and not a weight.

Now the `get_input_embeddings` returns a `WordEmbeddings` layer that has a `word_embeddings` attribute. If you look at the Bert model for example, the layer `TFBertLMPredictionHead` takes a `WordEmbeddings` layer as `input_embeddings` and use the `WordEmbeddings.word_embeddings` attribute into the `tf.matmul`.

> Now the `get_input_embeddings` returns a `WordEmbeddings` layer that has a `word_embeddings` attribute. If you look at the Bert model for example, the layer `TFBertLMPredictionHead` takes a `WordEmbeddings` layer as `input_embeddings` and use the `WordEmbeddings.word_embeddings` attribute into the `tf.matmul`.

So this part confuses me. Why name `word_embeddings` the weights inside the `WordEmbeddings`? It causes so much headache when reading the code afterward as we keep seeing some  `word_embeddings` attributes which might either be an embedding layer or a weight.

Also, how does the new organization not screw up pretrained weights? From what I understand, the old `world_embeddings` in the `BertEmbeddings` layer used to be a weight and now it's a layer with an added `world_embeddings` attribute?

> So this part confuses me. Why name word_embeddings the weights inside the WordEmbeddings? It causes so much headache when reading the code afterward as we keep seeing some word_embeddings attributes which might either be an embedding layer or a weight.

I agree it is confusing, if you prefer it can be called `weight` such as in `TFSharedEmbeddings` I think it would be a more suitable name. This renaming will make easier the kind of checking (from the incoming PR on ebd resizing)
```python
 def _get_word_embedding_weight(self, embedding_layer):
        if hasattr(embedding_layer, "word_embeddings"):
            return embedding_layer.word_embeddings
        elif hasattr(embedding_layer, "weight"):
            return embedding_layer.weight
        elif hasattr(embedding_layer, "decoder"):
            return embedding_layer.decoder
        else:
            # Here we build the word embeddings weights if not exists.
            # And then we retry to get the attribute once built.
            self(self.dummy_inputs)
            if hasattr(embedding_layer, "word_embeddings"):
                return embedding_layer.word_embeddings
            elif hasattr(embedding_layer, "weight"):
                return embedding_layer.weight
            elif hasattr(embedding_layer, "decoder"):
                return embedding_layer.decoder
            else:
                return None
```
No more `word_embeddings` or `weight`, only `weight`. What do you think?

> Also, how does the new organization not screw up pretrained weights? From what I understand, the old world_embeddings in the BertEmbeddings layer used to be a weight and now it's a layer with an added world_embeddings attribute?

This is because before we where using a [name score](https://www.tensorflow.org/api_docs/python/tf/name_scope) and not anymore in this PR. Let's say that defining a name scope or creating a layer represents the same thing. In both cases the weight is named `'tf_bert_model/bert/embeddings/word_embeddings/weight:0'` until now the `word_embeddings` part of the naming was because the embeddings was created in the context of `tf.name_scope("word_embeddings"):` , in this PR it has the same name but because of the name of the new `WordEmbeddings` layer.

Yes, having only "weight" makes more sense to me, and it would make the code easier to read. Thanks for explaining why the name of the weight doesn't change for loading!

I found another advantage of these new embedding computation. It allows our models to be compiled in XLA_GPU and XLA_TPU which was not the case before. Small proof test on a machine with a GPU:
```python
from transformers import TFBertModel
import tensorflow as tf

model = TFBertModel.from_pretrained("bert-base-cased")

@tf.function(experimental_compile=True)
def run():
    return model(model.dummy_inputs)

outputs = run()
```
On master fails with:
```
tensorflow.python.framework.errors_impl.InvalidArgumentError: Trying to access resource _AnonymousVar4 located in device /job:localhost/replica:0/task:0/device:CPU:0 from device /job:localhost/replica:0/task:0/device:GPU:0 [Op:__inference_run_4637]
```

On this PR works as expected. The reason is because the `tf.keras.layers.Embeddings` layers are initialized when the model is instanciated instead of being initialized at build time.

Now, each model has its own `WordEmbedding`, `TokenTypeEmbeddings` and `PositionEmbedding` layer in the model file decorated with the comment `#Copied from...` and the `words_embeddings` weights have been renamed into `weight` to make it more understandable and aligned with the name in `TFSharedEmbeddings`.

> LGTM in general. One thing I'm not 100% sure about is whether we really need to add keras layers like tf.keras.layers.Add() if we start doing this for the embeddings now, I'm wondering if we should do the same for all residual connections in the self-attention blocks

In the absolute, yes we should. In an ideal world, everytime TF proposes a function/layer for doing something we should use it, as it is part of the optimization process. I know and I understand that it might seems confusing and starts to diverge with what PT looks like.