TFBertForMaskedLM won't reload from saved checkpoint, shape mismatch issue
Author: martingajek
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.5.1-4.7
- Platform: Debian GNU/Linux 10 (buster)
- Python version: 3.9.2
- PyTorch version (GPU?): N/A
- Tensorflow version (GPU?): 2.5.0
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help
@Rocketknight1, @LysandreJik, @sgugger

## Information

Model I am using: TFBertForMaskedLM

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [X] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [ X] my own task or dataset: (give details below)

I believe this issue also affects the official TFTrainer implementation
as the checkpoint restore snippet was adapted from it.

## To reproduce

Steps to reproduce the behavior:

1. Generate Masked Batch
2. initialize TF Model and assign CheckpointManager
3. Save model checkpoint
4. initialize new TF Model and assign CheckpointManager
5. restore from checkpoint


```
import numpy as np
from transformers import AutoTokenizer, TFAutoModelForMaskedLM, AutoConfig, TFAutoModelForCausalLM
import tensorflow as tf

random_sentences = ["You'll see the rainbow bridge after it rains cats and dogs.",
"They looked up at the sky and saw a million stars.",
"The bullet pierced the window shattering it before missing Danny's head by mere millimeters.",
"He was willing to find the depths of the rabbit hole in order to be with her."]

tok = AutoTokenizer.from_pretrained('bert-base-uncased')
input_ids = tok.batch_encode_plus(random_sentences,return_tensors='np',padding=True)['input_ids']

#Create masked tokens as labels
labels = np.ones_like(input_ids)*-100
mask = (np.random.uniform(size=input_ids.shape)<=0.2) & (input_ids != 0)
labels[mask]=tok.mask_token_id

batch= {'input_ids':tf.convert_to_tensor(input_ids),
        'labels':tf.convert_to_tensor(labels)}

"""## Run model and save checkpoint"""

model = TFAutoModelForMaskedLM.from_pretrained('bert-base-uncased')
checkpoint = tf.train.Checkpoint(model=model)
model.ckpt_manager = tf.train.CheckpointManager(checkpoint, './', max_to_keep=1)
out = model(**batch)
print(out.loss.numpy())
model.ckpt_manager.save()

"""## Re-Initialize from config alone an load existing checkpoint"""

cfg = AutoConfig.from_pretrained('bert-base-uncased')
model2 = TFAutoModelForMaskedLM.from_config(cfg)
checkpoint2 = tf.train.Checkpoint(model=model2)
model2.ckpt_manager = tf.train.CheckpointManager(checkpoint2, './', max_to_keep=1)
latest_ckpt = tf.train.latest_checkpoint('./')
status = checkpoint2.restore(latest_ckpt)
status.assert_existing_objects_matched()

out = model2(**batch)
print(out.loss.numpy())
```

## Expected behavior

Expect to fully restore from checkpoint

## Current Behavior, error output

```---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-12-5ec2de12ee44> in <module>()
----> 1 out = model2(**batch)
      2 out.loss

19 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py in set_shape(self, shape)
   1238       raise ValueError(
   1239           "Tensor's shape %s is not compatible with supplied shape %s" %
-> 1240           (self.shape, shape))
   1241 
   1242   # Methods not supported / implemented for Eager Tensors.

ValueError: Tensor's shape (512, 768) is not compatible with supplied shape [2, 768]
```

## Link to colab
https://colab.research.google.com/drive/12pwo4WSueOT523hh1INw5J_SLpkK0IgB?usp=sharing



Comments:
This is an odd one - can you check if the problem still occurs when you use `model.save_pretrained()`? Just pass a path to that method to save, then load the model using `TFAutoModelForMaskedLM.from_pretrained()` with the same path.

Yeah `model.save_pretrained()` followed by `TFAutoModelForMaskedLM.from_pretrained()` works fine 

@Rocketknight1  I think that this is due to the fact that the variable for the token_type_embeddings clashes with the position_embeddings for having the same name: if I assign a different name to position_embeddings here: https://github.com/huggingface/transformers/blob/2e5dbdf2db4599a6694d0974575a70f9bc3c978e/src/transformers/models/bert/modeling_tf_bert.py#L164
Say go from "embeddings" to "pos_embeddings", the issue disappears, it's weird because I would expect the name_scope to take precedence but apparently not. I imagine that if I were to proceed with that name change, that might cause issues during model conversion to the BertModel pytorch implementation.

Hey, thank you for that very helpful bit of diagnostic info! That links this with #11202,  another issue we have caused by the same underlying problem. This is helpful because I'll probably need to make some breaking changes to fix that issue, and the fact that it's causing multiple downstream problems will increase the urgency there.

Cool! Glad that was helpful, thanks for looking into the issue.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.