src/transformers/models/vit/modeling_flax_vit.py/FlaxViTEmbeddings/setup
class FlaxViTEmbeddings:    def setup(self):
        self.cls_token = self.param("cls_token", nn.initializers.zeros, (1, 1, self.config.hidden_size))
        self.patch_embeddings = FlaxPatchEmbeddings(self.config, dtype=self.dtype)
        num_patches = self.patch_embeddings.num_patches
        self.position_embeddings = self.param(
            "position_embeddings", nn.initializers.zeros, (1, num_patches + 1, self.config.hidden_size)
        )
        self.dropout = nn.Dropout(rate=self.config.hidden_dropout_prob)