src/transformers/models/esm/modeling_tf_esm.py/TFRotaryEmbedding/__init__
class TFRotaryEmbedding:    def __init__(self, dim: int, name=None):
        super().__init__(name=name)
        # Matt: The PyTorch version of this layer does a lot of work to cache values, but we just rely on TF compilation
        # and/or XLA to sort out constants like that. It actually may not seem like this layer needs to be stateful at
        # all when we benefit from TF compilation, but it does. The reason is that self.inv_freq is a buffer in the
        # original implementation, but all the shared ESM checkpoints were trained with fp16 params. This means that
        # the inv_freq tensor was stored as a float16, and we need to replicate those lower-precision values or our
        # models give different outputs from the original.
        self.dim = dim