src/transformers/models/beit/modeling_flax_beit.py/FlaxBeitPatchEmbeddings/__call__
class FlaxBeitPatchEmbeddings:    def __call__(self, pixel_values):
        embeddings = self.projection(pixel_values)
        batch_size, _, _, channels = embeddings.shape
        return jnp.reshape(embeddings, (batch_size, -1, channels))