src/transformers/models/esm/modeling_tf_esm.py/TFRotaryEmbedding/call
class TFRotaryEmbedding:    def call(self, q: tf.Tensor, k: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        cos_emb, sin_emb = self._compute_cos_sin(k, seq_dimension=-2)

        return (
            apply_rotary_pos_emb(q, cos_emb, sin_emb),
            apply_rotary_pos_emb(k, cos_emb, sin_emb),
        )