bugs in flax_dinov2 when batch size is not 1
Author: Artimisu
Pull request: no
Labels: name: Flax

name: bug

### System Info

4.46

### Who can help?

Flax: @sanchit-gandhi flax: @sanchit-gandhi @MHRDYN7 

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

my code :

from transformers import AutoImageProcessor, FlaxDinov2Model
from PIL import Image 
import requests
import jax
import jax.numpy as jnp
import random
model = FlaxDinov2Model.from_pretrained("dinov2_base")
key = jax.random.PRNGKey(0)
inputs_np = {'pixel_values':jax.random.normal(key, (2, 3, 224, 224))}
outputs = model(**inputs_np)
last_hidden_states = outputs.last_hidden_state



when batch size is not 1, the bug is : 
<img width="1032" alt="image" src="https://github.com/user-attachments/assets/5c0ad809-cabc-4963-96b8-f69cb57567a0">

I reviewed the source code and found that in the function FlaxDinov2Embeddings/interpolate_pos_encoding within the file transformers/models/dinov2/modeling_flax_dinov2.py, the shape of class_pos_embed does not match the shape of patch_pos_embed. 

Specifically, the first dimension of class_pos_embed is 1, while the first dimension of patch_pos_embed corresponds to the batch size, so they cannot be concatenated.

<img width="865" alt="image" src="https://github.com/user-attachments/assets/42b49046-9925-4db7-81ff-429350dd2b7a">




 I would like to know how I can fix this bug and use dinov2 with a dataset batch size greater than 1.

### Expected behavior

I would like to know how I can fix this bug and use dinov2 with a dataset batch size greater than 1.

Comments:
@Artimisu I'll take a look at the issue. As far as I remember the [tests](https://github.com/huggingface/transformers/blob/main/tests/models/dinov2/test_modeling_flax_dinov2.py) did try a batch size of more than 1. If you come up with a fix, please feel free to open a pr.

Thank you for your response. I noticed that in your tests, when setting the batch size to a value greater than 1, the test_model function runs smoothly. However, the test_model_from_pretrained function encounters issues, as I mentioned earlier. 

This is because, in the test_model function, num_patches = num_positions. so the interpolate_pos_encoding function returns position_embeddings directly. 

<img width="753" alt="image" src="https://github.com/user-attachments/assets/ba879097-c939-43d7-a64a-dd1c64ca8e3c">

However, when using the pretrained model, it executes following code, encountering the mismatch between the first dimension of class_pos_embed and the first dimension of patch_pos_embed. 

As @Swastik-Swarup-Dash suggested, I can expand class_pos_embed to resolve the issue. It would be great if this bug could be addressed in the next official version of the code. Thank you again for your help.

I modified the code as follows, and it now runs smoothly:

patch_pos_embed = jnp.transpose(patch_pos_embed, (0, 2, 3, 1)).reshape((1, -1, dim))
patch_pos_embed_expanded = jnp.repeat(patch_pos_embed, hidden_states.shape[0], axis=0)
class_pos_embed_expanded = jnp.repeat(class_pos_embed[jnp.newaxis, :], hidden_states.shape[0], axis=0)
return jnp.concatenate((class_pos_embed_expanded, patch_pos_embed_expanded), axis=1)

It’s important to expand not only class_pos_embed but also patch_pos_embed. If patch_pos_embed is not expanded, an error will be thrown in the next line because the dimensions of the embeddings and the function’s return value will not match：

embeddings = embeddings + self.interpolate_pos_encoding(
            self.config, embeddings, height, width, self.position_embeddings
        )