Unable to use set_weights() and reload model later
Author: echan5
Pull request: no

### System Info

- `transformers` version: 4.38.2
- Platform: Linux-6.1.82-x86_64-with-glibc2.31
- Python version: 3.11.9
- Huggingface_hub version: 0.21.4
- Safetensors version: 0.4.2
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): not installed (NA)
- Tensorflow version (GPU?): 2.14.1 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help?

@gante @Rocketknight1

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

What is the proper way to train a new model with custom (rather than the default randomly initialized) weights?

I'd like to train a new (`TFLayoutLM`) model but want to load the embeddings layers with weights extracted from another pretrained model. However, after saving the new model, and trying to reload that saved model from file, that loaded model does not have the expected weights (rather, they're randomly initialized weights).

```
import numpy as np
from transformers import LayoutLMConfig, TFLayoutLMModel
 
configuration = LayoutLMConfig()
model = TFLayoutLMModel(configuration)

new_weights = [np.random.random(size=x.shape) for x in model.layoutlm.embeddings.weights]

model.build() 
model.layoutlm.embeddings.set_weights(new_weights)

model.save_pretrained('mysillymodel')
newmodel = TFLayoutLMModel.from_pretrained('mysillymodel')

# the weights are NOT as expected
np.testing.assert_equal(newmodel.layoutlm.embeddings.get_weights(), new_weights)
```

Reloading the saved model comes with the warning `Some layers from the model checkpoint at mysillymodel were not used when initializing TFLayoutLMModel: ['embeddings/token_type_embeddings/embeddings:0', 'encoder/layer_._2/attention/self/value/kernel:0'...You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference` (this warning shows even if I do train the model before saving out, which indicates it's not loading the trained weights but randomly initializing).

I'm not sure if the issue has something to do with having to call `build()` first before being able to set the weights, since if I don't set custom weights (or more generally, if I don't call `build()` even if not setting custom weights), then the weights are as expected.

Without `build()` call:
```
import numpy as np
from transformers import LayoutLMConfig, TFLayoutLMModel
 
configuration = LayoutLMConfig()
model = TFLayoutLMModel(configuration)

input_ids = np.array([[ 101, 7592, 2088,  102]])
bbox = np.array([[[0, 0, 0, 0], [ 637,  773,  693,  782], [ 698,  773,  733,  782], [1000, 1000, 1000, 1000]]])
preds = model.predict((input_ids, bbox))

expected_weights = model.layoutlm.embeddings.get_weights()
model.save_pretrained('mysillymodel')

newmodel = TFLayoutLMModel.from_pretrained('mysillymodel')

# the weights are as expected
np.testing.assert_equal(newmodel.layoutlm.embeddings.get_weights(), expected_weights)
```

Adding `build()` call:
```
import numpy as np
from transformers import LayoutLMConfig, TFLayoutLMModel
 
configuration = LayoutLMConfig()
model = TFLayoutLMModel(configuration)
model.build()  # add build

input_ids = np.array([[ 101, 7592, 2088,  102]])
bbox = np.array([[[0, 0, 0, 0], [ 637,  773,  693,  782], [ 698,  773,  733,  782], [1000, 1000, 1000, 1000]]])
preds = model.predict((input_ids, bbox))

expected_weights = model.layoutlm.embeddings.get_weights()
model.save_pretrained('mysillymodel')

newmodel = TFLayoutLMModel.from_pretrained('mysillymodel')

# the weights are NOT as expected
np.testing.assert_equal(newmodel.layoutlm.embeddings.get_weights(), expected_weights)
```

### Expected behavior

I'd expect the weights of the loaded model to be the same as that of the model that got saved out (i.e. the `np.testing.assert_equal` line in the example code above should not raise an AssertionError.

Comments:
@echan5 I have a suspicion about what's going on here - can you replace `model.build()` with `model.build_in_name_scope()` and tell me if it works after that?

Yes, that seems to fix it. Thank you

Yes, the cause of the problem is TF name scoping, as it often is with weight loading issues! The core issue is:

- There are two ways that `build()` can be called - directly, or implicitly when the model is called. If the model has not been built, `model.__call__()` will call `model.build()` before it calls `model.call()`
- If it's called implicitly by `__call__()`, then it will already be in a name scope. If it's called manually, it will not.
- Therefore, the weights of TF models can be mismatched if `build()` is called manually! We added `build_in_name_scope()` to get around this, but it's still annoying to have a separate method.
- Ideally, we'd autodetect if we're in a call scope in `build()` and enter a name scope if we aren't, which would resolve any discrepancies, but this can be quite messy - the Keras methods for identifying call scopes are private and tend to move around with each update, so it's very annoying to make this work for all supported version of TF/Keras. If you (or anyone reading this) can come up with a better solution, we'd welcome a PR!

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/main/CONTRIBUTING.md) are likely to be ignored.