src/transformers/models/esm/modeling_tf_esm.py/TFRotaryEmbedding/_compute_cos_sin
class TFRotaryEmbedding:    def _compute_cos_sin(self, x, seq_dimension=2):
        seq_len = tf.shape(x)[seq_dimension]

        t = tf.range(seq_len, dtype=self.inv_freq.dtype)
        freqs = tf.einsum("i, j -> ij", t, self.inv_freq)  # Outer multiplication
        emb = tf.concat((freqs, freqs), axis=-1)[None, None, :, :]

        return tf.cos(emb), tf.sin(emb)