Port IDEFICS to tensorflow
Author: a8nova
Pull request: yes

# What does this PR do?
This PR ports IDEFICS to tensorflow 
<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes # (issue)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
@VictorSanh just an fyi - I am hoping this model is TF portable.. 

also ccing @Rocketknight1 

🤗👀

Hey @a8nova! I'm the TF maintainer around here, and this sounds like a great idea. Feel free to ping me if you encounter any difficulties. In general when porting code to TF you will first want to do the 'boilerplate'. Take a look at #22970 to see how this works! All of these steps are fairly quick:

1) Copy the `modeling_idefics.py` file to `modeling_tf_idefics.py`
2) Rename all the classes e.g. `IdeficsForXXX -> TFIdeficsForXXX`
3) Add the relevant imports in `src/transformers/models/idefics/__init__.py` and `/src/transformers/__init__.py`
4) Add the TF classes to `modeling_tf_auto.py` 
5) Add the TF classes in `docs/source/en/model_doc/`

After this, you can start actually porting the code in `modeling_tf_idefics.py`. In general, you can just replace `torch.xxx()` ops with `tf.xxx()`, and layers from `torch.nn` with layers from `tf.keras.layers`. When creating layers, subclass from `tf.keras.layers.Layer` instead of `nn.Module`. Many of the ops are exactly identical, or only have minor name changes. There are a few key changes to watch out for, though:

6) The TF `forward()` method is called `call()`. NOT `__call__()`!
7) TF layers usually don't have an `input_dim` argument. This value is inferred when they are built with their first input.
8) When creating a layer (any class that is subclassed from `tf.keras.layers.Layer`, which includes TF built-in layers like `Dense`), pass the attribute name as the `name` argument. This ensures the TF weights layout lines up with the PyTorch state dict, like so:
```python
self.layer_norm1 = tf.keras.layers.LayerNormalization(name="layer_norm1")
```
9) In-place modifications of tensors are prohibited in TF (and JAX!). Most neural net code doesn't do this anyway, because it creates problems in backprop. If you need to do it, you can use something like `tf.where()` or `tf.scatter_nd_update()` to create a new tensor with the updates instead. This can be tricky, let me know if you need help!
10) For... reasons... the base stem of the model is contained in `TFIdeficsBaseLayer` and `TFIdeficsModel` is just a wrapper around this. Derived classes all use `TFIdeficsBaseLayer` and don't create any `TFIdeficsModel` layers. This is different from Torch models, where the model stem is contained in `IdeficsModel` and the derived classes use it as a layer.

We've actually had some success using GPT-4 to draft a port of the modeling code, so let me know if you'd like me to do that and add the GPT port of `modeling_tf_idefics.py` to this PR as a starting point!

After all this, the last step is to make changes to any processing code and tests needed to also support the TF model. It's a long list, but it's doable!

Thank you @Rocketknight1 for the detailed guide.  I have 1,2 & 3 done already, i just updated the PR. I will continue to work on the rest.

Regarding the GPT-4 generated draft, I already started doing some of the work, if you think the generated draft is easier to port to TF, please add it here and I can continue working from that ( I saw a comment about "auto-translation" in modeling_tf_sam.py and I was wondering about the details.. :)

A few questions:
1. What about all the `torch.nn` code inside `perceiver.py` and `vision.py`, do they also need TF porting? (my goal is to port inference code first, if this isn't needed for inference, then maybe i can come back to it)
2. For model_tf_auto.py, what is [this code](https://github.com/huggingface/transformers/blob/c030fc891395d11249046e36b9e0219685b33399/src/transformers/models/auto/modeling_tf_auto.py#L437) doing? It is not clear to me what to add the TF idefics versions, since i don't understand that file

Thanks for all the help!

Hi @a8nova, to answer the questions:

1) You'll probably need to convert those files too - the `IdeficsVisionTransformer` in `vision.py` seems to be a core part of the model. You might be able to skip `perceiver.py` initially, as it's only used in some model configs, but we should probably do it somewhere as part of this PR!
2) That code is there because SAM was the first `ForMaskGeneration` model we added in TF. For Idefics it's easier, because you're not adding a whole new "category" of model like that. If you look in `modeling_auto.py` you'll see there are two lines for `IdeficsModel` and `IdeficsForVisionText2Text` - all you need to do is make copies of those lines in `modeling_tf_auto.py` for `TFIdeficsModel` and `TFIdeficsForVisionText2Text`.

I'll work on generating GPT-4 translations for you, and post them here when they're available! Since you've already started on `modeling_tf_idefics.py` I won't overwrite it. You can just copy pieces from it when you need to.

@a8nova I added the three files with `_autotranslate.py` endings! Note that there are likely to be issues (e.g. forgetting the `name` kwarg when initializing layers even though I told it to)

Thank you @Rocketknight1, yes that makes sense, taking a closer look, idefics is 2 pre-trained models combined together, so vision.py is a core part.

I will take a look at the auto-translated files!

Hello @Rocketknight1 - Sorry I went MIA, I have been occupied with my 9-5. I just made some progress. I have the tests running, I am running into a weird error, I will attach a file below, any ideas?
[error.txt](https://github.com/huggingface/transformers/files/13373847/error.txt)


Also, regarding:
`"For... reasons... the base stem of the model is contained in TFIdeficsBaseLayer and TFIdeficsModel is just a wrapper around this. Derived classes all use TFIdeficsBaseLayer and don't create any TFIdeficsModel layers. This is different from Torch models, where the model stem is contained in IdeficsModel and the derived classes use it as a layer."
`
how come I don't see similar stuff for other TF implementations? Is this specific to IDEFICS? 

Hi @a8nova!

Firstly the error: The problem is that TF models don't let you assign to `self.layers`, because TensorFlow reserves that as a special keyword. What you should do is replace `self.layers` with something else, maybe `self.modules` or `self.decoder_layers` or something. However, you should **keep the name kwarg as** `layers_{i}`. We match TF layers to PT layers when doing weight cross-loading using the name attributes of layers, not the actual name like `self.layers`, so as long as you keep the argument the same then weight cross-loading should work.

Secondly, regarding `TFIdeficsBaseLayer`, that was actually just a typo on my part - it's actually called `TFIdeficsMainLayer`! If you check any of our other TF models, you'll see they have a `MainLayer` class like `TFBertMainLayer` that contains the model stem.

Thank you @Rocketknight1! 

Hi @Rocketknight1, I have a few questions:

1. For processing_idefics.py, how do you suggest I handle both pytorch and tf? right now I just have it hacked in my view to only have TF stuff (just to unblock me). 
2. I am getting this weird error from [freeze_model](https://github.com/huggingface/transformers/blob/b42fe29af93c70a1b3555af1b87aab32ed822f41/src/transformers/models/idefics/modeling_tf_idefics.py#L255) . I am doing something wrong but not sure, the error is `AttributeError: 'TFIdeficsRMSNorm' object has no attribute 'layers'` (full stacktrace attached). Any ideas?
3. There is also a "ALL_LAYERNORM_LAYERS" in the pytorch code, I added it in this commit https://github.com/huggingface/transformers/pull/26870/commits/7e0a35119b4d7a6284d04d8c543fba1b29e573c9,  does this look right to you?

Thanks in advance!

[error.txt](https://github.com/huggingface/transformers/files/13431615/error.txt)


Hi @a8nova, let's see...

For 1, we usually add an argument like `return_tensors` which can take values like `tf`, `pt`, etc. You can take a look at e.g. `models/sam/processing_sam.py` for an example - the main thing is that you should guard any imports of `tf` or `torch` behind something like `is_torch_available()` to make sure that the code doesn't crash for people who only have one of them installed!

For 2, I think the problem there is that `freeze_text_layers` iterates over multiple layers that include the normalization layer `self.norm`:
```python
def freeze_text_layers(self, module_exceptions=[]):
    for module in [self.decoder_layers, self.norm]:
        freeze_model(module, module_exceptions=module_exceptions)
```
As a result, it tries to iterate over `self.norm.layers`, which doesn't exist because `self.norm` is just a `LayerNormalization` layer, not a model with additional sub-layers. I'll suggest a change in `freeze_model` that should help.

For 3, `ALL_LAYERNORM_LAYERS` is a value mainly used by the HuggingFace `Trainer`, which is only used for PyTorch. When training with TF, we just use Keras instead to get the same functionality. You can just skip/remove it!

Thank you @Rocketknight1 for the detailed explanation and help! Moving on to other errors..

Sorry to bother you again @Rocketknight1, Do i need to worry about gradient checkpointing in the TF training tests? 

I am asking because for TFIdeficsForVisionText2Text there is `test_training` and `test_training_gradient_checkpointing` and they call `model.train()` and `model.gradient_checkpointing_enable()` which fail with `AtributeError`.

I found gradient_checkpointing_enable() inside modeling_utils.py, i don't see one inside modeling_tf_utils.py, can you guide me if my observation is correct and if I need to add all the gradient_checkpointing_* routines in modeling_tf_utils.py? 

Thank you!


@a8nova No, you can skip gradient checkpointing in TF ports!

Hi @Rocketknight1, how do i run the integration tests? For example `IdeficsModelIntegrationTest`.

`test_training` for tf is failing due to `model.train()`. i see it is defined in Trainer class in files trainer.py and trainer_tf.py. I don't think trainer_tf.py is used anywhere or is it? how do you suggest I resolve this? Thanks!

Hi @a8nova, you're right - we used to have a `TFTrainer` class but it's now deprecated. We recommend just training our TF models using the Keras API like `model.fit()`, and that means methods like `model.train()` do not really exist for them. I wrote a blogpost about this [here](https://huggingface.co/blog/tensorflow-philosophy) if you're curious, but it's mostly not that relevant to this PR!

Anyway, as a result of this the integration tests for TF models can be quite different from the integration tests for PT models - I'd recommend copying tests from another TF model instead, rather than trying to exactly copy the PT tests. You may also have to drop some tests entirely - the full IDEFICS tests use 4-bit quantization in PT, which isn't supported for TF, and as a result our CI machines may not be able to run it in TF at all!

Thanks @Rocketknight1. I have a few follow up questions. There are some test failures I don't understand.

1. For `image_processing_idefics.py`, I made changes in [this commit ](https://github.com/huggingface/transformers/pull/26870/commits/0a9fc014d431be78a5c4c2ab11e96e2ff41c47bf) to pass `return_tensors` to preprocess. I am testing it via the integration tests, I will share diff I needed to get it to run below. 

Does the integration test make sense?

```/content/transformers# git diff tests/models/idefics/test_modeling_tf_idefics.py
diff --git a/tests/models/idefics/test_modeling_tf_idefics.py b/tests/models/idefics/test_modeling_tf_idefics.py
index f9bcec579..50eee25d6 100644
--- a/tests/models/idefics/test_modeling_tf_idefics.py
+++ b/tests/models/idefics/test_modeling_tf_idefics.py
@@ -454,7 +454,7 @@ class TFIdeficsForVisionText2TextTest(TFIdeficsModelTest, unittest.TestCase):
 
 @require_tf
 @require_vision
-class IdeficsModelIntegrationTest(TestCasePlus):
+class TFIdeficsModelIntegrationTest(TestCasePlus):
     @cached_property
     def default_processor(self):
         return (
@@ -463,8 +463,6 @@ class IdeficsModelIntegrationTest(TestCasePlus):
             else None
         )
 
-    @require_bitsandbytes
-    @slow
     def test_inference_natural_language_visual_reasoning(self):
         cat_image_path = self.tests_dir / "fixtures/tests_samples/COCO/000000039769.png"
         cats_image_obj = Image.open(cat_image_path)  # 2 cats
@@ -494,9 +492,7 @@ class IdeficsModelIntegrationTest(TestCasePlus):
             load_in_4bit=True,
             bnb_4bit_compute_dtype="float16",
         )
-        model = IdeficsForVisionText2Text.from_pretrained(
-            "HuggingFaceM4/idefics-9b", quantization_config=quantization_config, device_map="auto"
-        )
+        model = TFIdeficsForVisionText2Text.from_pretrained("HuggingFaceM4/idefics-9b")
         processor = self.default_processor
         inputs = processor(prompts, return_tensors="tf")
         generated_ids = model.generate(**inputs, max_length=100)
```

How do i pass `return_tensors`? I adopted other code where `return_tensors` for image processing code is passed to the preprocess code only and not sure how to pass it here.


2. I have attached an error log for another failure I get. This is coming from TFIdeficsDecoupledLinear:
`E               assert <transformers.models.idefics.modeling_tf_idefics.TFIdeficsDecoupledLinear object at 0x7d4a9ac41540> is None`, 

full error in file attached


There are a few other errors I don't understand but some look related to this. Thanks for the help @Rocketknight1 !

Hi @a8nova, sorry for the delay!

Firstly, for `return_tensors`, generally our processors handle it like this:

```python
def __init__(self, return_tensors=None):
    self.return_tensors = return_tensors

def __call__(self, return_tensors=None):
    if return_tensors = None:
        return_tensors = self.return_tensors
```

In other words, you can supply it as an argument to `__call__`, but if you don't supply that argument, then it uses the default value that's set in the `__init__`.

Also, I took a look in the attached file but I can't see the error - can you double-check the contents?

whoops I attached the wrong one, i have attached the correct one below. thank you @Rocketknight1! 
Edit: i can't seem to repro that error now, I will look into it over the weekend again

No probs - let me know if it recurs!

Thank you @Rocketknight1!. There is another test failure I could use your help on. For `test_resize_token_embeddings`, the input `pixel_values` is `CHW` without the batches N, which code is responsible for resizing this kind of input? Right now it crashes when it tries to build the dummy weights, I have attached full error. The vision_tf.py code is still a little bugy. In the `TFIdeficsVisionEmbeddings` forward pass, the first thing I do is change the `pixel_values` to channels last format so it runs on CPU. what do I do when I get this kind of input? Thank you!

[test_resize_tokens_embeddings_error.txt](https://github.com/huggingface/transformers/files/13624772/test_resize_tokens_embeddings_error.txt)


Hi @a8nova, the models aren't intended to run with `CHW` input - they should always receive `NCHW`! The `test_resize_token_embeddings` test is mostly designed to work with text models - you might need to override for IDEFICS, but it's also not a critical test, so you can just skip it instead!

Also, the thing you did where you transpose to `NHWC` / `channels_last` is normal for TF models, because TF convolutions don't work on CPU in `channels_first` mode. Be careful that you transpose back after convolutional layers/blocks, though, or else the dimensions in the rest of the code will be wrong!

Hi @Rocketknight1, Happy holidays! I have a few more questions. I have less test failures now but I can't wait to get this model working end-to-end. I am using the integration test as a start for fully testing the tensorflow implmentation.

**EDIT**: Update on 1 (Jan 7th): I have figured out the issue, it was due to a bad reshape, I am able to run the model end-to-end now using tiny-random-idefics. I will run it with idefics-9b next! (I still have some follow up questions coming up as there is some weirdness i don't understand)

~~1. For the integration test, I have( This is from `TFIdeficsModelIntegrationTest` slightly modified as below, I understand this test will fail with `tiny-random-idefics` but I am using it to help me flush other bugs)~~
```
        model = TFIdeficsForVisionText2Text.from_pretrained("HuggingFaceM4/tiny-random-idefics", from_pt=True)
        processor = self.default_processor
        inputs = processor(prompts, return_tensors="tf")
        generated_ids = model.generate(**inputs, max_length=100)
        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
```
~~Right now this fails when calling from_pretrained, it goes into the forward pass and has an invalid size for pixel_values, this goes back to what I asked last time where I said the input is missing batches for some test but it looks like the problem is from my code actually. Right now what I am observing is that `processing_idefics.py::__call__` returns the correct size for `pixel_values` from the call to BatchFeatures() where `pixel_values` is a 5d tensor of [2,2,3,30,30] but the first forward pass inside modeling_idefics and vision_tf.py have `pixel_values` of [3,30,30].
Do you have suggestions how this might be?~~

2. I tried converting the pytorch weights to tf but failed with some error, do i need to get the tensorflow implementation working before I can convert the weights?

3. Is `dummy_inputs` still needed for tf implementations? Like [modeling_tf_bert.py#L916](https://github.com/huggingface/transformers/blob/4fd32a1f499e45f009c2c0dea4d81c321cba7e02/src/transformers/models/bert/modeling_tf_bert.py#L916)



Hi @a8nova, sorry for the Christmas-related delay. Huge congratulations on getting the tiny-random model working though, that indicates you're probably quite close to getting the whole thing working!

Firstly, I'd suggest that **you probably do not need to convert the weights**. The reason for this is that our `idefics` checkpoints all have `safetensors` weights already, and those weights can be loaded by any framework - TF, PyTorch or JAX. Once the TF code is working, you should just be able to load those repos with `from_pretrained` with no additional changes to the repo.

Secondly, dummy inputs are much less important than they used to be, since we're moving to adding explicit `build()` methods on all of our TF models. That's something we'll have to add to this PR as well, but let's leave it to the end - I can help with it, and it shouldn't be too painful.

Finally, I notice a lot of the test failures are caused by the TF test runner trying to `import torch` - the reason for this is that you're importing `ModelOutput` from `modeling_outputs.py`, but this is a torch file - try importing it from `modeling_tf_outputs` instead!

Still, I think this PR is getting close now - I should be able to reply much more quickly now the holidays are over, so please let me know if you encounter any other difficulties!

Also @alshenkute it seems like the issue there is that the model is stored in `bfloat16`, which is supported in PyTorch on GPU and CPU, but maybe not on M1. I'm not too familiar with developing on M1 though - maybe load with `torch_dtype=torch.float32`?

Got it, thank you @Rocketknight1! let me try that. ( i accidentally replied from another email, sorry for confusion, i deleted it)


Hi @Rocketknight1 - I am still having issues loading the idefics-9b model, Main issue right now is in the call to `TFIdeficsForVisionText2Text.from_pretrained`, `image_attention_mask` is None in TFIdeficsMainLayer. It looks like `image_attention_mask` is populated by `prepare_inputs_for_generation` which isn't called for the `from_pretrained`. I have a hack but it is ugly and doesn't feel correct.

For pytorch, the `from_pretrained` doesn't seem to invoke the inference code while loading the model weights, so i don't see same issue..

What am i missing here? If none of the stuff I am saying above makes sense  then there is a chance I have a bug elsewhere or I am misunderstanding things

By the way, it looks like the OOM is killing my process and I am looking into why that is but I need to understand how to handle the issue above with the call to `from_pretrained` 

Thanks in advance!

Hi @a8nova - I think I have a suspicion of what's going on now, let me explain!

I made a PR just before the holidays at #27794. This PR replaces our old approach to building TF models, where we used dummy inputs and ran an inference step to build weights, with a new approach where we build weights with proper `build()` methods. We did this because there were several issues with the dummy inputs, including poor performance, and I suspect the issue you're encountering is related to this.

I think the solution here is twofold:

1) Rebase this PR on the most recent `transformers` version (so you get the new `build()` PR)
2) Add `build()` methods to IDEFICS.

Since this is a very new PR that was added while your PR was open, I realize that it's a bit awkward. Let me know if you want me to handle one or both of those steps!

Thank you @Rocketknight1! Makes lots of sense.. 

I will give it a try, I just rebased with upstream, I will work on 2 next!

@a8nova - if it helps, I have an automated script I used in that PR that adds `build()` methods. Let me know if you want me to run it for you here!

@Rocketknight1 - Actually yes that would be great. Thank you @Rocketknight1 ! 

It also looks like we will need to do it to 2 files for IDEFICS, `modeling_tf_idefics.py` and `vision_tf.py`

Once you add the `build()` methods, I will re-run with idefics-9b!

Hi @a8nova, they're added! There may still be some issues - if you encounter errors later that are triggered inside a `build()` method, or that complain that a layer isn't getting the shape it's expecting, that's a sign of an issue with the automatically created methods, usually that the script inferred the wrong shape. Let me know if those start appearing and I'll work on fixes, but usually it's just a case of copying the input dim from the relevant PyTorch layer to the `build()` call to that layer in TF.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26870). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Thank you @Rocketknight1, I will pick this back up over the weekend!

Hi @Rocketknight1 - I quickly ran it & I am getting below error, not sure if it is issues with my code in general or due to the `build()` methods that were just added? Thanks!

```

Traceback (most recent call last):
  File "/Users/alazarshenkute/Projects/transformers-2024/transformers/src/transformers/models/idefics/idefics.py", line 73, in <module>
    run_idefics("tf")
  File "/Users/alazarshenkute/Projects/transformers-2024/transformers/src/transformers/models/idefics/idefics.py", line 36, in run_idefics
    model = TFIdeficsForVisionText2Text.from_pretrained(checkpoint,
  File "/Users/alazarshenkute/Projects/transformers-2024/transformers/src/transformers/modeling_tf_utils.py", line 2914, in from_pretrained
    return load_pytorch_checkpoint_in_tf2_model(
  File "/Users/alazarshenkute/Projects/transformers-2024/transformers/src/transformers/modeling_tf_pytorch_utils.py", line 195, in load_pytorch_checkpoint_in_tf2_model
    return load_pytorch_weights_in_tf2_model(
  File "/Users/alazarshenkute/Projects/transformers-2024/transformers/src/transformers/modeling_tf_pytorch_utils.py", line 236, in load_pytorch_weights_in_tf2_model
    return load_pytorch_state_dict_in_tf2_model(
  File "/Users/alazarshenkute/Projects/transformers-2024/transformers/src/transformers/modeling_tf_pytorch_utils.py", line 269, in load_pytorch_state_dict_in_tf2_model
    tf_model(tf_inputs, training=False)  # Make sure model is built
  File "/usr/local/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/Users/alazarshenkute/Projects/transformers-2024/transformers/src/transformers/modeling_tf_utils.py", line 425, in run_call_with_unpacked_inputs
    return func(self, **unpacked_inputs)
  File "/Users/alazarshenkute/Projects/transformers-2024/transformers/src/transformers/models/idefics/modeling_tf_idefics.py", line 1644, in call
    outputs = self.model(
  File "/Users/alazarshenkute/Projects/transformers-2024/transformers/src/transformers/models/idefics/modeling_tf_idefics.py", line 1466, in build
    self.vision_model.build(None)
  File "/Users/alazarshenkute/Projects/transformers-2024/transformers/src/transformers/models/idefics/vision_tf.py", line 555, in build
    self.embeddings.build(None)
  File "/Users/alazarshenkute/Projects/transformers-2024/transformers/src/transformers/models/idefics/vision_tf.py", line 166, in build
    self.patch_embedding.build(None)
TypeError: Exception encountered when calling layer 'tf_idefics_for_vision_text2_text' (type TFIdeficsForVisionText2Text).

'NoneType' object is not subscriptable

Call arguments received by layer 'tf_idefics_for_vision_text2_text' (type TFIdeficsForVisionText2Text):
  • input_ids={'input_ids': 'tf.Tensor(shape=(1, 2), dtype=int32)', 'attention_mask': 'tf.Tensor(shape=(1, 2), dtype=int32)', 'pixel_values': 'tf.Tensor(shape=(1, 3, 30, 30), dtype=float32)'}
  • attention_mask=None
  • position_ids=None
  • past_key_values=None
  • inputs_embeds=None
  • pixel_values=None
  • image_encoder_embeddings=None
  • perceiver_embeddings=None
  • image_attention_mask=None
  • labels=None
  • use_cache=None
  • output_attentions=None
  • output_hidden_states=None
  • interpolate_pos_encoding=False
  • return_dict=None
  • training=False

```

Hi @a8nova, these errors usually mean that one of the shapes the script created was wrong - in this case, the line `self.patch_embedding.build(None)` is failing because `None` is not an acceptable shape for that layer.

To fix errors like this, find the place where the layer is created in the Pytorch file, and see what the input dim for the layer is there, and then copy that to the `build()` method - so for example, if the input_dim in PyTorch is `self.num_channels`, then you'd change that line to `self.patch_embedding.build(self.num_channels)`. Keep doing that until you stop seeing errors - if you get stuck or you can't figure them out, let me know and I can help fix the issues!

Thank you @Rocketknight1 . I just fixed those errors. when running idefics-9b I noticed below warning in the logs:

`Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFIdeficsForVisionText2Text:
Some weights or buffers of the TF 2.0 model TFIdeficsForVisionText2Text were not initialized from the PyTorch model and are newly initialized:
`
the list for non initialized weights is very long, it looks like almost everything...

For the tf weight initialization, i am passing the attribute name as the name argument so it aligns with the pytorch weights, is there anything else i am missing? 

**EDIT jan-20:** Looks like I have narrowed it down, there were a lot of missing `name=` in tf weight initialization

Hi @a8nova, it's looking quite good now! I have one more request - if you get a chance, can you quickly fix the other imports that are causing the CI to fail because it tries to `import torch`? Usually this is just stuff like importing from a torch file like `modeling_utils` instead of `modeling_tf_utils`. Once you do that, the CI should clean up a lot, and it'll be easier for me and other reviewers to see what's left to be done without needing to download the PR branch and run it locally. I think we're getting quite close!

Hi @Rocketknight1 - Gotcha. I just tried fixing it in [this commit](https://github.com/huggingface/transformers/pull/26870/commits/5ae3ce08abfae7a773bbbbdbbbfb9de30b3f7ec3).

I am not sure why i don't see the import issue when running tests locally. how do i repro above CI failure locally? Can you tell me what to run in my view please. ( I looked at the pipeline but not able to repro).

Also, there are still a few issues in the port:
1. There are a few hacks i have locally that I don't know how to properly fix yet ( I will share details for all the hacks so you can help me resolve)
2. Score outputs SxS for pytorch vs tf doesn't match. In the commit I made today I fixed this issue for vision_tf.py
3. There are still about 20 test failures when running test_modeling_tf_idefics.py
4. I see non-deterministic behavior. For same runs, I see the model giving different outputs. I fixed this inside vision_tf.py when I moved initializing class_embedding into build(), the other fix needed was being able to load the pytorch state dict in tf
5. And model still doesn't work end to end, I am hoping once I have (2) done where scores match for SxS, TF port will start working!


Hi @a8nova, the reason you're probably not seeing those failures in your local instance is because you have `torch` installed there. The CI runs TF tests in a TF-only environment and torch tests in a torch-only environment, unless the tests are specifically marked as cross-tests.

As for the other stuff, I'll investigate!

Hi @a8nova - I did a pass on TFIDEFICS - the torch imports issue should be resolved as well as the type hint stuff and several other problems. There are still some remaining test failures - I suspect if we focus on the ones that throw shape errors first, that will hopefully resolve a lot of the numerical issues too. The CI should also now only show 'real' issues, now that the torch/type hint stuff isn't clogging it up.

See if you can trace any of the remaining shape issues, and let me know if you get stuck and I can dig in again!

thanks a ton @Rocketknight1 for making those fixes! I will look into the remaining issues

Hi @Rocketknight1! - Looks like I have addressed the shape issues, the non-determinstic behavior is also now gone, most likely due to one of your commits where you added the weight initialization inside`build()`. 

I will now focus on the SxS mismatch, the scores start to diverge after first forward pass and I am going to trace that now!

Awesome! That kind of stuff is usually just a small issue on one or two lines, so I think we're getting quite close now.

Hi @Rocketknight1 - a few questions:

1. I remember you said don't forget to transpose back to channels-first, where would be the right place to transpose back? I was thinking at the [end of the forward pass, transposing back the output patch_embeds](https://github.com/a8nova/transformers/blob/8881ac051dbec16dd8934a5288f9bfc1a5a67b27/src/transformers/models/idefics/vision_tf.py#L156) but the dimensions without the transpose back for rest of code looks correct... can you please confirm?
2. I have fixed one issue with the mismatch, epsilon wasn't being passed to a few LayerNormalization in the perceiver code. Now it looks like mismatch starts after a few iterations after the keras call [tf.keras.layers.Attention](https://github.com/a8nova/transformers/blob/8881ac051dbec16dd8934a5288f9bfc1a5a67b27/src/transformers/models/idefics/modeling_tf_idefics.py#L728) in the attention layer. just wondering if this is expected since the attention output after that call is very different and not a little off..
3. Just an fyi, i am still running this using tiny-random, i will run it using idefics-9b to see just how far off the actual output is (not the scores but text output). Right now when trying to load idefics-9b, OOM is killing me on Colab pro and my box (i noticed one huge drawback to cross-loading pytorch weights in tf is it needs/uses a lot of RAM because it loads pytorch state_dict then converts to tf), are there any other free high RAM machines I can use to quickly test?


Hi @a8nova! To answer the questions:

1) I would say you should transpose back after each convolution, or after a run of convolutions. The reason is that convolutions are usually mixed in with other layers like `BatchNorm`, `LayerNorm`, `Dense`, etc. If you've copied the code from PyTorch, it will assume that layout is channels-first, and that means these other layers will be operating on dim `1`. So I'd say the rule should be:

- If there's a single convolution, transpose to channels-last before the convolution and to channels-last after the convolution
- If there are multiple convolutions in a row with no other layers in between, transpose to channels-last before the first convolution and to channels-last after the last convolution 

2) Ah, there is actually an issue in the code there, and quite an annoying one! It won't work to create a `tf.keras.layers.Attention` layer in the `call()`, because Keras layers have their own weights - they project the inputs rather than just performing attention. In the original PyTorch code, this was `nn.functional.multi_head_attention`, and we'll need a similar function for TF. Let me work on that and I'll ping you when it's ready!

Hi @a8nova, I added a proper `scaled_dot_product_attention` function to `tf_utils` that should match the behaviour of torch's `nn.functional.scaled_dot_product_attention` up to numerical error, and replaced the call to `tf.keras.layers.Attention` with it. Hopefully the outputs should be a lot closer now - let me know if there are other issues remaining!

Thanks a ton @Rocketknight1!

Hi @Rocketknight1! the outputs are looking much better now, there are a few minor fixes I made that i will commit soon.

So the hidden_states in the MainLayer now matches after the first full forward pass. Now the logits have a mis-match (we are super close to getting this fully working 🤞) and I have narrowed it down to [TFIdeficsDecoupledLinear](https://github.com/a8nova/transformers/blob/1b51e8cf6d49e0f9c3d18bbdb734f463718029ac/src/transformers/models/idefics/modeling_tf_idefics.py#L416) 

`output = tf.linalg.matmul(input, self.weight)`, 

the [pytorch equivalent is](https://github.com/a8nova/transformers/blob/1b51e8cf6d49e0f9c3d18bbdb734f463718029ac/src/transformers/models/idefics/modeling_idefics.py#L428) 
 
 `output = F.linear(input, self.weight, self.bias)` 
 
do you see any issue with tf `output = tf.linalg.matmul(inputs, self.weight)` or that layer in general? I think this issue is similar to the `nn.functional.scaled_dot_product_attention`  issue you helped me fix above

Hi @a8nova, yeah, there is an issue. The key (which GPT-4 missed) is that the Torch layer is inheriting from `nn.Linear`. This means that sets a lot of attributes like `bias` itself in the superclass code. Our ported layer just inherits from `layers.Layer`, and so doesn't do any special stuff like that. As a result, I think it just sets `self.bias` to `True`. This, I suspect, will either cause an error, or will cause the bias to be uniformly set to `1` if the boolean value is cast to an integer!

However, that's kind of irrelevant because `bias=False` is forced in the code and so we never actually use it. As a result, I think the TF implementation should be fine. Can you check:

1) How big is the discrepancy between the outputs for TF and PyTorch for that layer?
2) Is that discrepancy equal to 1? Or `self.bias` for either layer?

Other than that, `tf.linalg.matmul()` should give the same output as `F.Linear()`

Hi @Rocketknight1 - thanks for the explanation. To answer your questions:
1. Discrepancy is all over the place for that layer output, but i did notice that right end of the matrix actually has same values (output below)... not sure what that means
2. Nope, `self.bias` is None for the models (idefics-9b & tiny-random)

I am also still getting a warning about some weights not being initialized for GatedCrossAttentionLayer with symbolic weight `'model.gated_cross_attn_layers.*.cross_attn.rotary_emb.inv_freq`, I tried fixing it in [this commit](https://github.com/huggingface/transformers/pull/26870/commits/d67cad586369916d98cc4c38dffdf52b6c1f9ca9) but I still see the warning, can you please take a look when you get a chance

I just ran latest code with idefics-9b, gives gibberish output.. I will keep digging..


```
=======Torch Decoupled Linear Output =========
 [[[-13.9692, -20.5610,   5.2010,  ...,  10.5567,  -2.0014,   4.9280],
 [-14.5836, -23.4861,   4.3161,  ...,   0.5876,  12.0983,   4.8461], 
 [ -9.2021, -13.5381,   3.2375,  ...,  18.9763,  -5.1458,   4.2513], 
 ======= Tf Decoupled Lienar Output ==========
 [[[-3.5217528   0.40634608  1.2831203  ... 10.556676   -2.0014238, 4.9279876 ]
[-3.16846    -0.3917833   2.6666665  ...  0.58874965 12.098683,  4.8463597 ]
[-5.0174932   0.91173846  0.64570814 ... 18.977085   -5.145547, 4.2512937 ]

```

Hey @a8nova, I patched up the build methods - I suspect the missing tensors problem was layers not being built correctly. Hopefully that also makes the outputs correspond a bit more closely - let me know if it's still happening after this fix and I'll dig in further!

Hi @Rocketknight1, thank you for the quick patch! I do still see the missing weights for the gated cross attention layer:

```
Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFIdeficsForVisionText2Text: ['model.vision_model.embeddings.position_ids', 'model.gated_cross_attn_layers.0.cross_attn.rotary_emb.inv_freq', 'model.gated_cross_attn_layers.1.cross_attn.rotary_emb.inv_freq']
```
thanks for the help!

Hi @a8nova, I dived in here - looking at the PyTorch code, it seems like `inv_freq` shouldn't be saved at all! In PyTorch it's a buffer with `persistent=False`. This actually enables a very easy rewrite of the TF code to replace it with a constant - hang on, I'll push something!

Done - the missing weights error should be gone now, though unfortunately it still hasn't fixed the outputs. Let me know if you get stuck there!

thank you @Rocketknight1! I will continue digging into the mismatch 

Hi @Rocketknight1! Which model did you run this with? I still see the weights not being initialized 

I just ran the tests with pytest - which model are you loading that shows the missing weights?

@Rocketknight1 - I am using tiny-random using integration test, below are the warnings for missing weights
```
Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFIdeficsForVisionText2Text: ['model.layers.1.self_attn.rotary_emb.inv_freq', 'model.gated_cross_attn_layers.0.cross_attn.rotary_emb.inv_freq', 'model.layers.0.self_attn.rotary_emb.inv_freq', 'model.vision_model.embeddings.position_ids', 'model.gated_cross_attn_layers.1.cross_attn.rotary_emb.inv_freq']
```

Hi @Rocketknight1! I think i am blocked by above missing weights issue, i don't think i can continue debugging the mis-match until that is resolved. Do you want me to share a Colab where you can quickly repro above issue using tiny-random? Thank you!


Yes please!

Hello @Rocketknight1 - Here is a [quick colab](https://colab.research.google.com/drive/1irzh32ThWcmobI3kVBRw324URLD4QLwG#scrollTo=FCaBJTSzUfec). Since the script is loading tiny-random, you can run it on a T4

**Just a few notes:**

There is a small hack in `image_processing_idefics.py` where I hardcode return_tensors since it is not being propagated and I haven't had a chance to look into that issue yet, so this new branch I created just has that hack, other than that is similar to the branch i am using for this PR.

Also just an FYI, first time you run the colab with a new runtime, please uncomment the two lines `setup_` and comment them again after runtime is initialized.

```
# uncomment below 2 lines for new runtimes
#setup_colab()
#setup_dev_env()
#%env TRANSFORMERS_VERBOSITY=''
!python src/transformers/models/idefics/idefics.py
```

Please let me know if you run into any issues.



Hi @a8nova, I have an answer for you at last, but it's quite a confusing one. What happened here is that in PyTorch, some tensors like `inv_freq` shouldn't be saved at all because they're buffers with `persistent=False`. However, **there are some existing checkpoints where the values are present**. This was probably caused by those checkpoints being made on an older version of the code, before the tensors were correctly set to `persistent=False`.

Therefore, the error you're getting is spurious, but if you want it to go away, I'd suggest making your own `tiny-random-idefics` like this:

```python
from transformers import IdeficsForVisionText2Text

model = IdeficsForVisionText2Text.from_pretrained("HuggingFaceM4/tiny-random-idefics")
model.push_to_hub("tiny-random-idefics")  # Will push to your Hub account
```

Now that you've loaded and re-saved the model, the tensors will be saved correctly, and the `persistent=False` tensors will not be included. Now you can just do:

```
model = TFIdeficsForVisionText2Text.from_pretrained("your_username/tiny-random-idefics")
```
and everything should be resolved!

Hi @Rocketknight1, thanks a lot! I can confirm after loading my re-saved model those misleading warnings are gone! I will continue debugging the mis-match this weekend.

Hi @Rocketknight1 - I have the model fully working now, after [this fix](https://github.com/huggingface/transformers/pull/26870/commits/90711174f0f05a4f01922cf1b74c610b00a22cae).

I tested it with `idefics-9b-instruct` and `tiny-random` with many different prompts!

Some remaining tasks needed to wrap up this PR:
1. Fix any remaining hacks I have (e.g. return_tensors isn't propagated)
2. Fix all failing tests when running `py.test -vv tests/models/idefics/test_modeling_tf_idefics.py` and the CI test failures. For this, I will need your help so we can figure out best way to move forward with some of the failures. I will look into all failures and share here details
3. Code review

Please let me know if I am missing anything. I would like to wrap up this PR ASAP

@a8nova no, that's everything! We really appreciate the work you did on this one - IDEFICS is a large and complex model. We'll try to support you to get this PR in as quickly as possible! 

Hi @Rocketknight1 - Can you guide me on how i can fix the 3 CI failures `check_code_quality`, `check_repository_consistency`, and `tests_pr_documentation_tests`. Thank you!

Hi @a8nova, those are triggered by our code style tools. To fix them, `cd` to the `transformers` dir and run `pip install -e .[quality]` followed by `make fixup`. If you get an error about too many files in `make fixup`, then you can manually run `make style` and `make fix-copies` - that should fix everything!

Hi @Rocketknight1 - I need your help with below test failures and how I should resolve them:
1.  First 2 failures `ci/circleci: examples_torch ` and `ci/circleci: tests_exotic_models` are failing when trying to git clone, is this related to my PR?
2. Also, i am not sure how to fix `tests_pr_documentation_tests`
3. What should i do regarding test failure`TFIdeficsModelTest::test_resize_token_embeddings` ? Should i skip it for this model? 


Also, unrelated to this PR, I wanted to port the tts model vits to tf, should i create another PR and will you create a draft for me like you did for this PR? Also, do you have any other model architecture in mind that should be ported instead?

Thank you!

Hi @a8nova, the failing torch/exotic tests are an issue with our CI right now - it'll be fixed soon, at which point you should be able to rebase. As for the PR documentation tests, the issue is that our doctest runner is failing when the TF code tries to load sharded safetensors checkpoints. I'll work on a PR to enable that! In the meantime, you can either remove those doctests or point them to a smaller model somewhere.

Hi @a8nova, draft PR to support sharded safetensors is up at #29350. I'll let you know when it's merged!

Thanks a ton @Rocketknight1!

Hi @Rocketknight1  - For the test failure `tests_torch_and_tf `, i am not sure how to resolve it, do you have any suggestions? 

Also, would you say this is ready for code review?

Hi @a8nova, the failing tf-torch tests look like there's a missing tensor in the TF output. You can insert breakpoints in the tests to run them locally and see what's going on, but most likely the cause is that torch is returning an extra tensor (like `attentions` or something) that TF isn't.

If you want to try running the tests locally, you can install torch (CPU version is fine) and `pytest`, and then do `RUN_SLOW=1 RUN_PT_TF_CROSS_TESTS=1 pytest tests/models/idefics/test_modeling_tf_idefics.py`. If you do it this way, you can insert breakpoints and immediately see what exactly the failure is!

thank you @Rocketknight1. I fixed it in [this commit](https://github.com/huggingface/transformers/pull/26870/commits/1b97814e085da407c8b771e14a0a6bac1cd5bde0). Please let me know if this isn't the proper fix. thanks!

Hi @Rocketknight1 -  I had to touch non idefics code in [this commit](https://github.com/huggingface/transformers/pull/26870/commits/4e9b7b44b77335a7a52e70d9b7eb8c9703deee6b). This was to fix error:
```
======================================================================================================================= FAILURES ========================================================================================================================
________________________________________________________________________________________________________ TFIdeficsModelTest.test_keras_save_load ________________________________________________________________________________________________________

self = <tests.models.idefics.test_modeling_tf_idefics.TFIdeficsModelTest testMethod=test_keras_save_load>

    def test_keras_save_load(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
    
        tf_main_layer_classes = {
            module_member
            for model_class in self.all_model_classes
            for module in (import_module(model_class.__module__),)
            for module_member_name in dir(module)
            if module_member_name.endswith("MainLayer")
            # This condition is required, since `modeling_tf_clip.py` has 3 classes whose names end with `MainLayer`.
            and module_member_name[: -len("MainLayer")] == model_class.__name__[: -len("Model")]
            for module_member in (getattr(module, module_member_name),)
            if isinstance(module_member, type)
            and keras.layers.Layer in module_member.__bases__
            and getattr(module_member, "_keras_serializable", False)
        }
        for main_layer_class in tf_main_layer_classes:
            # T5MainLayer needs an embed_tokens parameter when called without the inputs_embeds parameter
            if "T5" in main_layer_class.__name__:
                # Take the same values than in TFT5ModelTester for this shared layer
                shared = TFSharedEmbeddings(99, 32, name="shared")
                config.use_cache = inputs_dict.pop("use_cache", None)
                main_layer = main_layer_class(config, embed_tokens=shared)
            else:
                main_layer = main_layer_class(config)
    
>           symbolic_inputs = {
                name: keras.Input(tensor.shape[1:], dtype=tensor.dtype)
                for name, tensor in inputs_dict.items()
                #if tf.is_tensor(tensor)
            }

tests/test_modeling_tf_common.py:382: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

.0 = <dict_itemiterator object at 0x15865b420>

    symbolic_inputs = {
>       name: keras.Input(tensor.shape[1:], dtype=tensor.dtype)
        for name, tensor in inputs_dict.items()
        #if tf.is_tensor(tensor)
    }
E   AttributeError: 'bool' object has no attribute 'shape'

tests/test_modeling_tf_common.py:383: AttributeError
__
```

I want to avoid modifying non IDEFICS code so i don't break something else. Let me know if this is a problem. thank you!

Hi @a8nova, that change should be fine - the test should have been compatible with boolean inputs in the first place!

Hello @Rocketknight1  - The last 8 test failures are getting harder to fix. I looked into a few and they mostly have to do with tensorflow model saving for IDEFICS. I have started overriding `test_compile_tf_model` but I am getting many different errors since IDEFICS is just not compatible with None values, and even after fixing those I am getting more errors. 

I have an idea: Do you think i can skip these failures and mark them as TODO, merge in this PR (after review) and come back to them in a few weeks? I will be occupied with something else next few weeks & I don't want to leave this PR as a draft.

 I have attached a file with the remaining 8 failures for you to see, actually they are less than 8 since there are duplicate tests from `TFIdeficsModelTest` & `TFIdeficsForVisionText2TextTest`, so fixing it for one will fix it for both.

What do you think? Are there any I can skip temporarily?

[idefics_failures.txt](https://github.com/huggingface/transformers/files/14822978/idefics_failures.txt)


@Rocketknight1 - I need your input on test `TFIdeficsModelTest::test_loss_computation`, it is failing at 

```
if "labels" in prepared_for_class:
                labels = prepared_for_class["labels"].numpy()
                if len(labels.shape) > 1 and labels.shape[1] != 1:
                    labels[0] = -100
                    prepared_for_class["labels"] = tf.convert_to_tensor(labels)
                    loss = model(model_input, **prepared_for_class)[0]
                    self.assertTrue(loss.shape.as_list() == expected_loss_size or loss.shape.as_list() == [1])
>                   self.assertTrue(not np.any(np.isnan(loss.numpy())))
E                   AssertionError: False is not true

tests/test_modeling_tf_common.py:1400: AssertionError

```

The problematic line seems to be `labels[0] = -100` where it it setting all the elements to `-100` and causes [the loss_mask](https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_tf_utils.py#L236) to be zero which leads to NaN, how should i address this? Thanks!

Hi @a8nova, firstly sorry for the delay! For `test_loss_computation`, you can either skip the test, or copy the base test to the IDEFICS test file, and change it so that it actually correctly tests the loss! We do skip the test for several other TF models, and it's not included in the Torch tests, so I think it's okay to skip (though I'd make sure we check loss somewhere in the TF-PT equivalence tests)

As for the other test failures, can you give me the full trace for the `OperatorNotAllowedInGraphError` error? That usually indicates some forbidden operation that is preventing the model from compiling, so we do want to fix that one if possible!

Thank you @Rocketknight1, i understand if you are busy!

Actually i resolved the `OperatorNotAllowedInGraphError` issue, it was due to a None value being passed and `_make_causal_mask` doesn't like that. I think i am getting close to getting the compilation working, let me look into it further and if i get stuck i will just bug you here again. thanks!

Hi @Rocketknight1 - A few questions:

1 - The test [test_save_load](https://app.circleci.com/pipelines/github/huggingface/transformers/89356/workflows/9683755a-eb72-4594-a551-8fae41fb6fcf/jobs/1164665) passes on my local machine but fails when running on the CI..? Not sure if there are so many parallel tests running it is causing this or it is pointing to a bug

2 - Regarding the test `test_compile_tf_model`, the `OperatorNotAllowedInGraphError` was caused by the shape checks in vision_tf and modeling_tf_idefics which I just commented out:

```
+++ b/src/transformers/models/idefics/vision_tf.py
@@ -219,11 +219,11 @@ class TFIdeficsVisionAttention(tf.keras.layers.Layer):
         src_len = shape_list(key_states)[1]
         attn_weights = tf.linalg.matmul(query_states, key_states, transpose_b=True)
 
-        if shape_list(attn_weights) != [bsz * self.num_heads, tgt_len, src_len]:
-            raise ValueError(
-                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
-                f" {shape_list(attn_weights)}"
-            )
+        #if shape_list(attn_weights) != [bsz * self.num_heads, tgt_len, src_len]:
+        #    raise ValueError(
+        #        f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
+        #        f" {shape_list(attn_weights)}"
+        #    )

```
After commenting out above, i was running into a few weird errors when invoking [functional_model.save(tmpdirname)](https://github.com/huggingface/transformers/pull/26870/commits/8611ff959ce6d764e26219e1ebfe6e860b11b4eb#diff-27e43f726686e50e385006d911e2adc0d54fb2ada8d414aa6c62df77da7fd1eaR805), inside the `test_compile_tf_model` I overrode for TF-IDEFICS

For one (full stack trace attached)  pixel_values being passed is a non-tensor type (look at *2 below), even after fixing that error with a hack, I see batches being passed to vision_tf is`None` after first forward pass by `functional_model.save(tmpdirname)`. 

When I dump out `pixel_values` in vision_tf.py::TFIdeficsVisionEmbeddings::call

```
*1 pixel_values before transpose  Tensor("pixel_values:0", shape=(1, 3, 30, 30), dtype=float32)
*2 pixel_values before transpose  {'pixel_values': <tf.Tensor 'pixel_values:0' shape=(None, 3, 30, 30) dtype=float32>}
*3 pixel_values before transpose  Tensor("pixel_values:0", shape=(None, 3, 30, 30), dtype=float32)

```
My theory is that during regular inference we only run vision encoder forward pass (vision_tf.py) once for each image input, so once we have the features which are the outputs of the vision encoder, we pass that to perceiver resampler which gives us the visual tokens, then we move on to all the layers in modeling_tf_idefics.py code (gated cross, decoder layers, etc..), importantly, we only run vision encoder (& perceiver) forward pass once for each batch. But when `functional_model.save(tmpdirname)` is invoked after the initial forward pass we again invoke the vision encoder with some invalid input and that won't work with this model definition.. Let me know what you think, thanks!

[test_compile_tf_model_error.txt](https://github.com/huggingface/transformers/files/14949267/test_compile_tf_model_error.txt)

@a8nova Regarding 1, that's quite tricky! I don't remember a case when that test failed on CI but passed locally. Let's leave it for now, we can come back to it if we've fixed everything else but it's still a problem.

For 2, you can either comment out shape checks like that, or replace them with asserts from `tf.debugging` which should hopefully not trigger the `OperatorNotAllowed` errors.

For the `test_compile_tf_model` issues, it looks to me like some code is failing when flexible shapes are used. Most of our tests pass real eager tensors to the model, but when you compile a functional model, symbolic tensors are used instead, and these tensors can have `None` dimensions in their shape. To debug this, I recommend:

1) Initialize the model as normal
2) Pass some dummy inputs in and make sure it works (you can just use model.dummy_inputs)
3) Replace the dummy inputs with `tf.keras.Input` with the same shape, then pass those to the model (this is basically equivalent to building a functional model)
4) Try setting some of the dimensions in your `tf.keras.Input` to `None`. See if you can identify which `None` dimensions cause the error, and track down the problem.

I encountered bugs like this in the past in TF conversions, and often the cause was that there was some kind of `tf.reshape()` or `tf.transpose()` or `flatten()` or something like that which 'blended' the batch dimension with another dimension, and TF couldn't figure things out. For example, torch code might do something like this:

```
bs = x.shape[0]
x = x.flatten(0, 1)
# do some stuff here
x = x.unflatten(bs, -1)
```

This works in Torch because it knows the exact batch size, and can use that to make the rest of the code work, but TF doesn't know the exact batch size when it's compiling with flexible shapes, and so passing a shape like `(batch_size, -1)` will fail because `batch_size` is `None`!


Don't panic, though - I think this port is really close to being finished, so I suspect you'll only have to rewrite a couple of reshapes/transposes to make it work!

Thank you for the breakdown @Rocketknight1! I was making wrong assumptions, I think I understand now that things should work with flexible shapes (None batch_shape)


Hi @Rocketknight1, `test_compile_tf_model` is now passing, fixed in [this commit](https://github.com/huggingface/transformers/pull/26870/commits/049e094497c588538ae015a386e3ff60660fef29).

The `test_pr_documentation` is only failing test now, I thought your sharded tensors commit was suppose to fix that failure, or is there anything I need to do like pushing weights to the hub?

Hi @a8nova, this looks great! That does seem like a bug with safetensors loading in the doctest - I'll investigate it tomorrow. After that, we should be very close to merging. Thanks for all the work you put into this!

I am excited to get this merged in :). This wouldn't have been possible without your patience and guidance last few months, so I can't thank you enough for that!

There is some minor cleanup I will be doing on this PR, and as a final step I will also run this PR on a A100 by loading the 9B model to make sure there are no remaining issues

Hi @a8nova, I investigated the issue with the doctests, and I think the issue is just that another rebase onto `transformers/main` is needed! The TF safetensors updates haven't made it into your PR branch yet.

However, I should mention that `idefics-9b` is probably too big for the doctest runner anyway! For doctests like that, you can replace the `>>>` at the start of each line with `>>`. This looks the same for users, but the doctest runner won't pick it up and try to run it.

Hi @Rocketknight1 -  Thanks for catching that, rebasing did fix it, all tests now pass. I have marked this PR as ready for review!

Cool! Will review today.

Thank you for the review @Rocketknight1, I will address all of them!

Hi @Rocketknight1 - I think I addressed everything, I also just rebased and all tests are passing!

Cool, will re-review today!

Thank you for catching that @Rocketknight1, the pytorch processing code now uses it's original code from main

LGTM now! Doing some final checks before we merge.

Hi @a8nova, I realized we were missing some `is_pt_tf_cross_test` markers on the cross-tests, and when I added them I've started seeing failures in the CI. However, I can't reproduce them locally, the tests work fine on my machine.

Can you check it out and see if you can reproduce the errors? These errors are usually caused by the model being incompletely built somehow, causing weights to not be present when it's loaded. 

thanks for catching those missing markers! I just tried to repro and it passes on my local machine. I also tried on a Colab T4 but the test passes. Does the CI run the cross tests on the CPU or GPU? 

I believe on GPU, but I can't see how that would cause this! It's a very unusual issue.