src/transformers/models/opt/modeling_flax_opt.py/FlaxOPTLearnedPositionalEmbedding/setup
class FlaxOPTLearnedPositionalEmbedding:    def setup(self):
        self.offset = 2
        self.embedding = self.param(
            "embedding", self.embedding_init, (self.num_embeddings + self.offset, self.features), self.param_dtype
        )