Added RoPE Offset + Test

[Artur-Galstyan]

This PR follows #799 and adds the possibility for an offset in RoPE.

I've rebased my changes on the latest main branch and added one more test.

For your convenience, here's some code to test it with MHA:

Toggle me

import functools
import equinox as eqx
import jax
import jax.numpy as jnp
from equinox.nn._attention import MultiheadAttention
from equinox.nn._embedding import RotaryPositionalEmbedding
from jaxtyping import Array, Float, Int


embedding_size = 32
max_seq_length = 8
seq_length = 4
num_heads = 2
query_size = 64


class TransformerBlock(eqx.Module):
    rope_embeddings: RotaryPositionalEmbedding
    mha_attention: MultiheadAttention

    def __init__(self, embedding_size, max_seq_length, num_heads, query_size):
        self.rope_embeddings = RotaryPositionalEmbedding(embedding_size, max_seq_length)
        self.mha_attention = MultiheadAttention(
            num_heads=num_heads, query_size=query_size, key=jax.random.key(0)
        )

    def __call__(self, query, key_, value, index):
        def process_heads(
            query_heads: Float[Array, "seq_length num_heads qk_size"],
            key_heads: Float[Array, "seq_length num_heads qk_size"],
            value_heads: Float[Array, "seq_length num_heads vo_size"],
            index: Int[Array, ""],
        ) -> tuple[
            Float[Array, "seq_length num_heads qk_size"],
            Float[Array, "seq_length num_heads qk_size"],
            Float[Array, "seq_length num_heads vo_size"],
        ]:
            # index is the autoregressive index of the current token
            rope_partial = functools.partial(self.rope_embeddings, offset=index)
            query_heads = jax.vmap(rope_partial, in_axes=1, out_axes=1)(query_heads)
            key_heads = jax.vmap(rope_partial, in_axes=1, out_axes=1)(key_heads)

            return query_heads, key_heads, value_heads

        x = self.mha_attention(
            query=query,
            key_=key_,
            value=value,
            process_heads=functools.partial(process_heads, index=index),
        )

        return x


transformer_block = TransformerBlock(
    embedding_size, max_seq_length, num_heads, query_size
)
transformer_block = eqx.filter_jit(transformer_block)


q = jnp.ones(shape=(seq_length, query_size))
k = jnp.ones(shape=(seq_length, query_size))
v = jnp.ones(shape=(seq_length, query_size))

out = transformer_block(q, k, v, 0)
out = transformer_block(q, k, v, 1)
out = transformer_block(q, k, v, 2)
out = transformer_block(q, k, v, 3)
out = transformer_block(q, k, v, 4)
out = transformer_block(q, k, v, 5)