class InfiniAttention problem

class InfiniAttention(nn.Module):
r"""Hopefully faithful implementation of the InfiniAttention mechanism.

Read the paper: ["Leave No Context Behind: Efficient Infinite Context Transformers
with Infini-attention"](https://arxiv.org/abs/2404.07143)

Args:
embedding_dim (int):
The dimension of the input embeddings.
attn_head_dim (int):
The dimension of each attention head used in the multi-head attention.
num_query_heads (int):
The number of query attention heads.
num_key_value_heads (int):
The number of key and value attention heads.
use_attn_linear_bias (bool):
Whether to use a bias in the linear layer after attention.
use_delta_update_rule (bool):
Whether to use the delta update rule mentioned "Section 2.1.2 Compressive Memory Update"
in the paper.
"""

def __init__(
self,
embedding_dim: int,
attn_head_dim: int,
num_query_heads: int,
num_key_value_heads: int,
use_attn_linear_bias: bool = False,
use_delta_update_rule: bool = False,
) -> None:
super().__init__()

self.embedding_dim = embedding_dim
self.attn_head_dim = attn_head_dim
self.query_dim = attn_head_dim * num_query_heads
self.key_value_dim = attn_head_dim * num_key_value_heads
self.num_query_heads = num_query_heads
self.num_key_value_heads = num_key_value_heads
self.use_delta_update_rule = use_delta_update_rule

self.q_proj = nn.Linear(embedding_dim, self.query_dim, bias=False)
self.k_proj = nn.Linear(embedding_dim, self.key_value_dim, bias=False)
self.v_proj = nn.Linear(embedding_dim, self.key_value_dim, bias=False)
self.attn = ScaledDotProductAttention(attn_head_dim)
self.linear = nn.Linear(
self.key_value_dim, embedding_dim, bias=use_attn_linear_bias
)
self.elu = nn.ELU()
self.sigmoid = nn.Sigmoid()

# beta is used in long-term context injection for each attention head
# [b? n 1 1] x [b n s v] => [b n s v]
self.beta = nn.Parameter(torch.randn((num_key_value_heads, 1, 1)))

# key: [batch_size, num_key_value_heads, seq_length, attn_head_dim]
# key_T: [batch_size, num_key_value_heads, attn_head_dim, seq_length]
# memory: dim(key_T) * dim(value)
# [batch_size, num_kv_heads, attn_head_dim, seq_length] * [batch_size, num_kv_heads, seq_length, attn_head_dim]
# => [batch_size, num_kv_heads, attn_head_dim, attn_head_dim]
memory = torch.zeros((num_key_value_heads, attn_head_dim, attn_head_dim))
z = torch.zeros((num_key_value_heads, self.attn_head_dim))
self.register_buffer("memory", memory)
self.register_buffer("z", z)

def forward(self, query: T, key: T, value: T, mask: Optional[T] = None) -> T:
# 1. Projection
q_proj: T = self.q_proj(query)
k_proj: T = self.k_proj(key)
v_proj: T = self.v_proj(value)

# 2. Split into attention heads and transpose to get expected shape
batch_size, seq_length, _ = q_proj.shape
q_proj = q_proj.view(batch_size, -1, self.num_query_heads, self.attn_head_dim)
k_proj = k_proj.view(
batch_size, -1, self.num_key_value_heads, self.attn_head_dim
)
v_proj = v_proj.view(
batch_size, -1, self.num_key_value_heads, self.attn_head_dim
)

q_proj = q_proj.transpose(1, 2)
k_proj = k_proj.transpose(1, 2)
v_proj = v_proj.transpose(1, 2)

# 2.1 Retrieve from memory
elu_q = self.elu(q_proj) + 1

# numerator: [b n_q s a] x [b? n_kv a a] => [b n??? s a]
# denominotor: [b n_q s e] x [b? n_kv e 1] => [b n??? s 1]
# TODO: n??? For now, num_query_heads must be equal to num_key_value_heads otherwise this will fail
# TODO: Implement GQA correctly
A_mem = torch.matmul(elu_q, self.memory) / (
torch.matmul(elu_q, self.z.unsqueeze(dim=-1)) + EPSILON
)

# 2.2 Memory update
elu_k: T = self.elu(k_proj) + 1
elu_k_T = elu_k.transpose(2, 3)

if self.use_delta_update_rule:
v_delta = torch.matmul(elu_k, self.memory) / (
torch.matmul(elu_k, self.z.unsqueeze(dim=-1)) + EPSILON
)
v = v_proj - v_delta
else:
v = v_proj

self.memory = self.memory + torch.matmul(elu_k_T, v)
self.z = self.z + elu_k.sum(dim=2)

# 3. SDPA
A_dot, context = self.attn(q_proj, k_proj, v_proj, mask)

# 3.1 Long-term context injection
beta = self.sigmoid(self.beta)
x = beta * A_mem + (1 - beta) * A_dot

# 4. Concat
x = x.transpose(1, 2).contiguous()
x = x.view(batch_size, -1, self.key_value_dim)

# 5. Linear
x = self.linear(x)

return x, context

A_dot, context = self.attn(q_proj, k_proj, v_proj, mask)
Is there any problem here? Why is this q_ proj, k_ proj, v_ proj instead of q_ proj [seq], k_ proj [seq] j, v_ projj [seq]?