Add TF ViT MAE
This PR adds the MAE [1] model in TensorFlow. It was developed by @arig23498 and myself. 

Fun facts about this PR:

* Probably the third pure vision model in TensorFlow in `transformers`.

**References**:

[1] [Masked Autoencoders Are Scalable Vision Learners](https://arxiv.org/abs/2111.06377)

**Update**

The PR is now ready for review. @gante @Rocketknight1 @sgugger 
I just have a quick look and left a comment.

I feel strange that there are style changes like `omega = 1.0 / 10000 ** omega # (D/2,)`. Do you have a previous version of `black`?

I thought you already updated the version (during your work on `TFConvNext`), via

```
pip install -e .[quality]
```
Maybe you were in a different virtual Python environment while working on this PR?<|||||>> I just have a quick look and left a comment.
> 
> I feel strange that there are style changes like `omega = 1.0 / 10000 ** omega # (D/2,)`. Do you have a previous version of `black`?
> 
> I thought you already updated the version (during your work on `TFConvNext`), via
> 
> ```
> pip install -e .[quality]
> ```
> 
> Maybe you were in a different virtual Python environment while working on this PR?

So, first I deactivate the current Python virtual environment and then run the installation and run `make style`? 

I think I should fetch upstream before that and rebase. <|||||>
> So, first I deactivate the current Python virtual environment and then run the installation and run `make style`?

Yeah, if your current venv is specific to your other work/projects, and you don't want to change its installed packages.
Maybe it would better if you create a new virtual environment, say, `venv-transformers-pr`, and switch to it.

> 
> I think I should fetch upstream before that and rebase.

You can try it. I always try to have (very) recent commit from master to work on a new PR. Hope the rebase is smooth in your case here.

<|||||>_The documentation is not available anymore as the PR was closed or merged._<|||||>Thank you for working on this model, @sayakpaul (and the patience for the randomness issue!).

I went through the parts (in the test files) involving the randomness and left a few comments.<|||||>@sgugger @ydshieh thank you for your positive comments. In the world of open-source, we couldn't have asked for me. So, thank you once again. 

We've addressed the PR comments. Let us know your thoughts. <|||||>On top of the code quality issue flagged above, the two following tests are failing:
```
FAILED tests/vit_mae/test_modeling_vit_mae.py::ViTMAEModelTest::test_pt_tf_model_equivalence
FAILED tests/vit_mae/test_modeling_tf_vit_mae.py::TFViTMAEModelTest::test_pt_tf_model_equivalence
```

Could you fix them so we can merge the PR? Thanks a lot!<|||||>> On top of the code quality issue flagged above, the two following tests are failing:
> 
> ```
> FAILED tests/vit_mae/test_modeling_vit_mae.py::ViTMAEModelTest::test_pt_tf_model_equivalence
> FAILED tests/vit_mae/test_modeling_tf_vit_mae.py::TFViTMAEModelTest::test_pt_tf_model_equivalence
> ```
> 
> Could you fix them so we can merge the PR? Thanks a lot!

I looked into it but couldn't figure out why they failed. Could you provide us with a starting point to get them fixed? <|||||>Here is the full stack trace:
```
============================= FAILURES SHORT STACK =============================
________________ TFViTMAEModelTest.test_pt_tf_model_equivalence ________________


tf_model = <transformers.models.vit_mae.modeling_tf_vit_mae.TFViTMAEModel object at 0x7f8bef271450>
pt_state_dict = OrderedDict([('embeddings.cls_token', tensor([[[ 0.0355,  0.0288,  0.0095,  0.0219, -0.0116,  0.0049,  0.0247,
       ...0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0.]))])
tf_inputs = {'pixel_values': <tf.Tensor: shape=(13, 3, 30, 30), dtype=float32, numpy=
array([[[[9.84787583e-01, 7.89411783e-01, 8....20547986e-01, 9.40254509e-01, ...,
          1.77839562e-01, 9.66387928e-01, 8.72015655e-01]]]],
      dtype=float32)>}
allow_missing_keys = False

    def load_pytorch_weights_in_tf2_model(tf_model, pt_state_dict, tf_inputs=None, allow_missing_keys=False):
        """Load pytorch state_dict in a TF 2.0 model."""
        try:
            import tensorflow as tf  # noqa: F401
            import torch  # noqa: F401
            from tensorflow.python.keras import backend as K
        except ImportError:
            logger.error(
                "Loading a PyTorch model in TensorFlow, requires both PyTorch and TensorFlow to be installed. Please see "
                "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions."
            )
            raise
    
        if tf_inputs is None:
            tf_inputs = tf_model.dummy_inputs
    
        if tf_inputs is not None:
            tf_model(tf_inputs, training=False)  # Make sure model is built
        # Adapt state dict - TODO remove this and update the AWS weights files instead
        # Convert old format to new format if needed from a PyTorch state_dict
        old_keys = []
        new_keys = []
        for key in pt_state_dict.keys():
            new_key = None
            if "gamma" in key:
                new_key = key.replace("gamma", "weight")
            if "beta" in key:
                new_key = key.replace("beta", "bias")
            if new_key:
                old_keys.append(key)
                new_keys.append(new_key)
        for old_key, new_key in zip(old_keys, new_keys):
            pt_state_dict[new_key] = pt_state_dict.pop(old_key)
    
        # Make sure we are able to load PyTorch base models as well as derived models (with heads)
        # TF models always have a prefix, some of PyTorch models (base ones) don't
        start_prefix_to_remove = ""
        if not any(s.startswith(tf_model.base_model_prefix) for s in pt_state_dict.keys()):
            start_prefix_to_remove = tf_model.base_model_prefix + "."
    
        symbolic_weights = tf_model.trainable_weights + tf_model.non_trainable_weights
        tf_loaded_numel = 0
        weight_value_tuples = []
        all_pytorch_weights = set(list(pt_state_dict.keys()))
        missing_keys = []
        for symbolic_weight in symbolic_weights:
            sw_name = symbolic_weight.name
            name, transpose = convert_tf_weight_name_to_pt_weight_name(
                sw_name, start_prefix_to_remove=start_prefix_to_remove, tf_weight_shape=symbolic_weight.shape
            )
    
            # Find associated numpy array in pytorch model state dict
            if name not in pt_state_dict:
                if allow_missing_keys:
                    missing_keys.append(name)
                    continue
                elif tf_model._keys_to_ignore_on_load_missing is not None:
                    # authorized missing keys don't have to be loaded
                    if any(re.search(pat, name) is not None for pat in tf_model._keys_to_ignore_on_load_missing):
                        continue
>               raise AttributeError(f"{name} not found in PyTorch model")
E               AttributeError: vit.embeddings.cls_token not found in PyTorch model

src/transformers/modeling_tf_pytorch_utils.py:198: AttributeError
_________________ ViTMAEModelTest.test_pt_tf_model_equivalence _________________


tf_model = <transformers.models.vit_mae.modeling_tf_vit_mae.TFViTMAEModel object at 0x7f8c0e80add0>
pt_state_dict = OrderedDict([('embeddings.cls_token', tensor([[[-0.0341, -0.0326,  0.0065, -0.0189,  0.0114,  0.0198,  0.0247,
       ...0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0.]))])
tf_inputs = {'pixel_values': <tf.Tensor: shape=(13, 3, 30, 30), dtype=float32, numpy=
array([[[[0.00428793, 0.00807637, 0.5928316 ...,
         [0.0279441 , 0.08120351, 0.30608034, ..., 0.777531  ,
          0.6238623 , 0.38510746]]]], dtype=float32)>}
allow_missing_keys = False

    def load_pytorch_weights_in_tf2_model(tf_model, pt_state_dict, tf_inputs=None, allow_missing_keys=False):
        """Load pytorch state_dict in a TF 2.0 model."""
        try:
            import tensorflow as tf  # noqa: F401
            import torch  # noqa: F401
            from tensorflow.python.keras import backend as K
        except ImportError:
            logger.error(
                "Loading a PyTorch model in TensorFlow, requires both PyTorch and TensorFlow to be installed. Please see "
                "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions."
            )
            raise
    
        if tf_inputs is None:
            tf_inputs = tf_model.dummy_inputs
    
        if tf_inputs is not None:
            tf_model(tf_inputs, training=False)  # Make sure model is built
        # Adapt state dict - TODO remove this and update the AWS weights files instead
        # Convert old format to new format if needed from a PyTorch state_dict
        old_keys = []
        new_keys = []
        for key in pt_state_dict.keys():
            new_key = None
            if "gamma" in key:
                new_key = key.replace("gamma", "weight")
            if "beta" in key:
                new_key = key.replace("beta", "bias")
            if new_key:
                old_keys.append(key)
                new_keys.append(new_key)
        for old_key, new_key in zip(old_keys, new_keys):
            pt_state_dict[new_key] = pt_state_dict.pop(old_key)
    
        # Make sure we are able to load PyTorch base models as well as derived models (with heads)
        # TF models always have a prefix, some of PyTorch models (base ones) don't
        start_prefix_to_remove = ""
        if not any(s.startswith(tf_model.base_model_prefix) for s in pt_state_dict.keys()):
            start_prefix_to_remove = tf_model.base_model_prefix + "."
    
        symbolic_weights = tf_model.trainable_weights + tf_model.non_trainable_weights
        tf_loaded_numel = 0
        weight_value_tuples = []
        all_pytorch_weights = set(list(pt_state_dict.keys()))
        missing_keys = []
        for symbolic_weight in symbolic_weights:
            sw_name = symbolic_weight.name
            name, transpose = convert_tf_weight_name_to_pt_weight_name(
                sw_name, start_prefix_to_remove=start_prefix_to_remove, tf_weight_shape=symbolic_weight.shape
            )
    
            # Find associated numpy array in pytorch model state dict
            if name not in pt_state_dict:
                if allow_missing_keys:
                    missing_keys.append(name)
                    continue
                elif tf_model._keys_to_ignore_on_load_missing is not None:
                    # authorized missing keys don't have to be loaded
                    if any(re.search(pat, name) is not None for pat in tf_model._keys_to_ignore_on_load_missing):
                        continue
>               raise AttributeError(f"{name} not found in PyTorch model")
E               AttributeError: vit.embeddings.cls_token not found in PyTorch model

src/transformers/modeling_tf_pytorch_utils.py:198: AttributeError
```
Looks like there is a weight present in the TF model but not the PyTorch one.<|||||>This does seem weird. I think if the TF model didn't have the right variable names as same as the PT one, it would not have loaded the PT weights back into the TF model, no?<|||||>@sgugger there seems to be something weird here. I ran the PT model like so:

```py
from transformers import ViTMAEForPreTraining
model = ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base")

param_states = model.state_dict()

for k in param_states:
	if "cls" in k:
		print(k)
```

It prints:

```
vit.embeddings.cls_token
```

But as per the error trace it says:

```
AttributeError: vit.embeddings.cls_token not found in PyTorch model
```
<|||||>Considering it was me working on this PT/TF equivalence test recently, I will try to take a look!
(despite the failure is more in `load_pytorch_weights_in_tf2_model` method)<|||||>@sayakpaul Could you try if the following work 🙏?

```
TFViTMAEModel.from_pretrained("facebook/vit-mae-base", from_pt=True)
```

OK, I saw it is tested indeed in your previous comment.<|||||>@sayakpaul 

change [here] from 
```
base_model_prefix = "vit_mae"
```
to
```
base_model_prefix = "vit"
```
will fix the loading issue, because that is the value used in the `ViTMAEModel`, see [here](https://github.com/huggingface/transformers/blob/2da1f03f93af37480fdfc23cbc870f5bf038c478/src/transformers/models/vit_mae/modeling_vit_mae.py#L585).

However I still get the large difference value between PT/TF
```
E   AssertionError: 3.805268 not less than or equal to 1e-05
```
So something still need to be fixed.

~~@sgugger @NielsRogge Do you think ~~

~~- `vit_mae` is a better value for `base_model_prefix`?~~
~~- if so, should we & are we able to fix this in PyTorch `ViTMAEModel`, considering there might be some (or many) downloads already?~~
~~- (I am completely fine if we don't change it because we can't)~~


Turns out `vit` is the final decision, see

https://github.com/huggingface/transformers/pull/16119/files (#16119)


So @sayakpaul , You need to change the base_model_prefix = "vit_mae" to "vit" in the TF model file. And we will try to see why we get large difference value.<|||||>@ydshieh thanks much for pointing us in the right direction. As for the assertion, it likely can be solved by fixing the `noise` argument like the way we did for the rest of the tests. 

If you could suggest where the `noise` argument needs to be added, I can look into it. 

**Update**:

Inside `check_pt_tf_models()` we need to pass `noise` to `pt_model` and `tf_model` to ensure the equivalence assertion passes. We can parameterize the `check_pt_tf_models()` method to allow `noise`. But since it's a core library component, let's discuss the best possible way to incorporate this change. 
<|||||>@sayakpaul Wow, super! `noise` is something we should keep in mind from now :-)

You can copy/paste the `test_pt_tf_model_equivalence` from `test_modeling_tf_common.py`, then add the `noise` part.
There are a few places we can remove (they only make sense in the common test), but if you are not sure which parts to remove, you can keep it as it is - I will make the change.<|||||>Thank you @ydshieh. Making the changes right away. <|||||>Seems like GitHub Actions is acting weird. But that's not what GitHub's status page says: https://www.githubstatus.com/.

Let's see if the current change can fix the equivalence test. <|||||>> Seems like GitHub Actions is acting weird. But that's not what GitHub's status page says: https://www.githubstatus.com/.
> 
> Let's see if the current change can fix the equivalence test.

I had exactly the same situation yesterday - 0 check after force push <|||||>I see. If possible it'd be great if you could review the changes I made so far. <|||||>@sgugger is it possible to trigger the workflows manually from your end? <|||||>Not from side, no. Can you do an empty commit (`git commit -m "Trigger tests" --allow-empty`)?<|||||>@sgugger seems like things are looking good. Let me know if the recent changes to fix equivalence look good to you. 

This is one last bit that I think you folks need to take care of: https://github.com/huggingface/transformers/pull/16255#discussion_r833968556. <|||||>> @sgugger seems like things are looking good. Let me know if the recent changes to fix equivalence look good to you.
> 
> This is one last bit that I think you folks need to take care of: [#16255 (comment)](https://github.com/huggingface/transformers/pull/16255#discussion_r833968556).

Glad to see the tests being all green!

@gante Before merge, please wait for me to remove some irrelevant parts in the overwritten PT/TF equivalence test, thanks! <|||||>I have pushed a commit to remove irrelevant blocks (that only make sense in the common test) + put pytorch inputs on the correct device (for GPU testing to run)<|||||>Thanks for your help, @ydshieh. The failing tests seem to be unrelated to the PR?<|||||>> Thanks for your help, @ydshieh. The failing tests seem to be unrelated to the PR?

Yeah, I think it is another model `facebook/xglm-564M` causing issue.
You can have a bit of rest for now 💯 💖

I re-run the failed test. ~~Let's see.~~ All green now<|||||>@gante according to https://github.com/huggingface/transformers/pull/16255/commits/debad62d47f5789f9890d3176713345e3a2506c0 there were no key additions or deletions from the code base. Do you know the possible reason for the tests to fail?<|||||>@gante @ydshieh 

Updates:

Repo consistency tells us that copied components have inconsistencies. Sorry, if this sounds foolish but we could not actually figure out those inconsistencies. Class name changes and other relevant changes in the comment do not count here I hope. To ensure the copying was indeed right we went ahead used a comparator tool to verify if that's the case. Here are the individual links to the comparison results for the components that are said to have copy inconsistencies:

* SelfAttention: https://www.diffchecker.com/m5VY19Xl
* SelfOutput: https://www.diffchecker.com/j6Pri9fN
* MAEAttention: https://www.diffchecker.com/5qtn9IRh
* Intermediate: https://www.diffchecker.com/hNy0Sg7c
* ViTOutput: https://www.diffchecker.com/CyK4zZUl
* Layer: https://www.diffchecker.com/BESXzSTD
* Encoder: https://www.diffchecker.com/BXHYohx1

Note that we did run `make fix-copies` within the environment as described in https://github.com/huggingface/transformers/pull/16255#discussion_r830432539 (refer to https://github.com/huggingface/transformers/pull/16255/commits/565ec4c1227d9577e0e4e7489880dad2d4c33d01) and it replaced `ViTMAEConfig` with `ViTConfig` which is wrong, I guess. 

We have also followed the copy comment format from the PT script ensuring they are well aligned. 

What else are we missing out here? <|||||>@ariG23498 @sayakpaul 

Probably my [previous comment](https://github.com/huggingface/transformers/pull/16255#discussion_r835623030) didn't explain things clear enough.

Let's take an example with:

```
# Copied from transformers.models.vit.modeling_tf_vit.TFViTSelfAttention with TFViT->TFViTMAE
class TFViTMAESelfAttention(tf.keras.layers.Layer):
    def __init__(self, config: ViTMAEConfig, **kwargs):
```

This `# copied from` will check if the block is a copy from `TFViTSelfAttention` after replaceing `TFViT` by `TFViTMAE` (in the RAM).
However, there is `ViTConfig` in the ViT file, and `ViTMAEConfig` in this block **(note: there is no TF prefix here)**, and the instruction `with TFViT->TFViTMAE` in the comment `# Copied from` won't apply to  `ViTConfig` (from the ViT file), and during the check, it is `ViTConfig` vs `ViTMAEConfig`, and therefore there is a difference.

In order to fix the issue, they are 2 potential options:

1. In `# Copied from`, change ` with TFViT->TFViTMAE` to ` with ViT->ViTMAE`
    - This will only work if the 2 blocks in ViT and ViTMAE are indeed the same after this replace!

2. In `# Copied from`, change ` with TFViT->TFViTMAE` to ` with TFViT->TFViTMAE, ViTConfig->ViTMAEConfig`
    - This is more verbose, but sometimes not really necessary

So I would suggest try option 1 first. If there are remaining places, we can use option 2 for that remaining places.
<|||||>For the remaining # copied from issue, 

```
# in ViT, layernorm is applied before self-attention
```
to
```
# in ViTMAE, layernorm is applied before self-attention
```
should work.
(This is also what has done in PyTorch ViTMAE)<|||||>It turns out that there is still 

```
# in ViT, layernorm is also applied after self-attention
```
to be changed ...I didn't check the whole block previously, sorry.

I don't always pull the latest commit from this PR. In general, it would be easier to detect if you run `make fix-copies` locally and see what have been changed, and you will get the idea where should be fixed.<|||||>Thanks for the prompt feedback @ydshieh

It was my bad to not check the entire code block.<|||||>Don't worry. 

And just a **potential** correction (not really important): my previous comment "it would be easier to detect ..." **might** be not True in this particular line: I don't think there will be a diff visible to this particular line, after running `make fix-copies`.

<|||||>All green ✅ 

Thanks @ydshieh for your prompt and valuable feedback!<|||||>About `shape_list() is a mess for graph mode. We have a plan to get rid of it, and just use .shape everywhere `

https://github.com/huggingface/transformers/pull/16255#discussion_r835525428

I was a bit worried it will break some testing at this moment, so I tested with `shape`. Indeed, the test `test_save_load` will give error at the line `model.save_pretrained(tmpdirname, saved_model=True)` - when calling `autograph_handler` inside. See the full log below.

## Full error log

```python
            with tempfile.TemporaryDirectory() as tmpdirname:
>               model.save_pretrained(tmpdirname, saved_model=True)

tests\vit_mae\test_modeling_tf_vit_mae.py:616: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
src\transformers\modeling_tf_utils.py:1418: in save_pretrained
    self.save(saved_model_dir, include_optimizer=False, signatures=self.serving)
..\..\..\..\miniconda3\envs\py39\lib\site-packages\keras\utils\traceback_utils.py:67: in error_handler
    raise e.with_traceback(filtered_tb) from None
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

args = ({'pixel_values': <tf.Tensor 'pixel_values:0' shape=(None, None, None, None) dtype=float32>},)
kwargs = {}

    def autograph_handler(*args, **kwargs):
      """Calls a converted version of original_func."""
      # TODO(mdan): Push this block higher in tf.function's call stack.
      try:
        return autograph.converted_call(
            original_func,
            args,
            kwargs,
            options=autograph.ConversionOptions(
                recursive=True,
                optional_features=autograph_options,
                user_requested=True,
            ))
      except Exception as e:  # pylint:disable=broad-except
        if hasattr(e, "ag_error_metadata"):
>         raise e.ag_error_metadata.to_exception(e)
E         ValueError: in user code:
E         
E             File "C:\Users\33611\Desktop\Projects\transformers-huggingface\transformers\src\transformers\models\vit_mae\modeling_tf_vit_mae.py", line 728, in serving  *
E                 return self.call(inputs)
E             File "C:\Users\33611\Desktop\Projects\transformers-huggingface\transformers\src\transformers\modeling_tf_utils.py", line 816, in run_call_with_unpacked_inputs  *
E                 return func(self, **unpacked_inputs)
E             File "C:\Users\33611\Desktop\Projects\transformers-huggingface\transformers\src\transformers\models\vit_mae\modeling_tf_vit_mae.py", line 1074, in call  *
E                 loss = self.forward_loss(pixel_values, logits, mask)
E             File "C:\Users\33611\Desktop\Projects\transformers-huggingface\transformers\src\transformers\models\vit_mae\modeling_tf_vit_mae.py", line 1007, in forward_loss  *
E                 target = self.patchify(imgs)
E             File "C:\Users\33611\Desktop\Projects\transformers-huggingface\transformers\src\transformers\models\vit_mae\modeling_tf_vit_mae.py", line 975, in patchify  *
E                 imgs = tf.cond(
E         
E             ValueError: Tried to convert 'x' to a tensor and failed. Error: None values not supported.

..\..\..\..\miniconda3\envs\py39\lib\site-packages\tensorflow\python\framework\func_graph.py:1129: ValueError
```<|||||>Changing `shape_list` to `shape` will also cause a test failing 

https://github.com/huggingface/transformers/blob/b320d87eceb369ea22d5cd73866499851cb2cca3/tests/utils/test_modeling_tf_core.py#L124-L135

with similar error

```python
E             in user code:
E             
E                 File "C:\Users\33611\Desktop\Projects\transformers-huggingface\transformers\src\transformers\modeling_tf_utils.py", line 816, in run_call_with_unpacked_inputs  *
E                     return func(self, **unpacked_inputs)
E                 File "C:\Users\33611\Desktop\Projects\transformers-huggingface\transformers\src\transformers\models\vit_mae\modeling_tf_vit_mae.py", line 1074, in call  *
E                     loss = self.forward_loss(pixel_values, logits, mask)
E                 File "C:\Users\33611\Desktop\Projects\transformers-huggingface\transformers\src\transformers\models\vit_mae\modeling_tf_vit_mae.py", line 1007, in forward_loss  *
E                     target = self.patchify(imgs)
E                 File "C:\Users\33611\Desktop\Projects\transformers-huggingface\transformers\src\transformers\models\vit_mae\modeling_tf_vit_mae.py", line 980, in patchify  *
E                     tf.debugging.assert_equal(imgs.shape[1], imgs.shape[2])
E             
E                 ValueError: None values not supported.
```


**(This test is currently only used for just a few core NLP models)**

**(This test will still fail with `shape_list` here, but the error seems coming from other lines rather than from shape_list itself)**<|||||>Thanks for investigating it, @ydshieh. Is there anything we can do in this PR to mitigate the problem? <|||||>> Thanks for investigating it, @ydshieh. Is there anything we can do in this PR to mitigate the problem?

Let's see what gante think.<|||||>> I was a bit worried it will break some testing at this moment, so I tested with shape. Indeed, the test test_save_load will give error at the line model.save_pretrained(tmpdirname, saved_model=True) - when calling autograph_handler inside. See the full log below.

This is why I used `shape_list()` there. The root cause here is that inside `patchify()` with `.shape` there will be a `None` in the batch size and `tf.random.uniform()` will error out for that. I could not think of any other workaround to mitigate the problem. 

For the second one (https://github.com/huggingface/transformers/pull/16255#issuecomment-1079785554) I am not too sure since you mentioned that the test is only applied for some core NLP models. 
<|||||>> For the second one ([#16255 (comment)](https://github.com/huggingface/transformers/pull/16255#issuecomment-1079785554)) I am not too sure since you mentioned that the test is only applied for some core NLP models.

My comments above are not about `we should change to shape and fix all the failing tests`.

It is more about a question (to gante) `if we should make this decision and change in this PR, considering some tests will fail.`
(And the second comment is included mainly for us (HF) not to forget there are some tests we want to run for all models ideally, but currently run only for a few models. It might be better for us (HF) not to make a big decision so quickly and see all tests being green for a particular model, but will fail some important tests that are not currently run due to some limitation).

<|||||>Now I understand. Appreciate the class clarification. <|||||>@sayakpaul `vit-mae-base` [TF weights are on the hub](https://huggingface.co/facebook/vit-mae-base/tree/main), the others will soon follow :) I think you can make the final changes now (remove the `from_pt`), to then merge the PR 💪 <|||||>@gante thank you! Changes made. <|||||>> @gante thank you! Changes made.

Can you confirm that the tests run with `RUN_SLOW=1`? Will merge after CI gets to green and I get your confirmation 🚀 <|||||>btw ignore the failure in `Add model like runner / Add new model like template tests (pull_request)`, it's being looked after<|||||>> > @gante thank you! Changes made.
> 
> Can you confirm that the tests run with `RUN_SLOW=1`? Will merge after CI gets to green and I get your confirmation 🚀

Yes, I did run the tests before pushing the changes. <|||||>@gante over to you to take the reigns.<|||||>CI green, slow tests were run, all TF weights on the hub 👉 merging.

Great work @sayakpaul @ariG23498 💪 