src/transformers/models/vit/modeling_flax_vit.py/FlaxPatchEmbeddings/__call__
class FlaxPatchEmbeddings:    def __call__(self, pixel_values):
        x = self.projection(pixel_values)
        batch_size, _, _, channels = x.shape
        return jnp.reshape(x, (batch_size, -1, channels))