transformer-lib 0.1.0__tar.gz

transformer_lib-0.1.0/PKG-INFO

@@ -0,0 +1,47 @@
+Metadata-Version: 2.4
+Name: transformer-lib
+Version: 0.1.0
+Summary: Attention mechanisms from Every_Thing_About_Transformer, as an importable library
+Author: Aniket Verma
+License: MIT
+Requires-Python: >=3.9
+Requires-Dist: torch>=2.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Provides-Extra: linear
+Requires-Dist: linear-attention-transformer; extra == 'linear'
+Description-Content-Type: text/markdown
+
+# transformer-lib
+
+Attention mechanisms from Every_Thing_About_Transformer, as an importable Python library.
+
+## Installation
+
+```bash
+pip install -e .
+```
+
+## Usage
+
+```python
+from transformer_lib import MultiHeadAttention
+import torch
+
+mha = MultiHeadAttention(embed_dim=512, num_heads=8)
+x = torch.randn(2, 10, 512)
+out = mha(x, x, x)
+```
+
+## Available Modules
+
+| Module | Description |
+|---|---|
+| `SelfAttention` | Basic scaled dot-product self-attention |
+| `MultiHeadAttention` | Multi-head attention with separate Q/K/V projections |
+| `CausalAttention` | Multi-head with lower-triangular causal mask |
+| `CrossAttention` | Encoder-decoder cross-attention with self-attention pre-step |
+| `GlobalAttention` | Additive (Bahdanau-style) alignment scoring |
+| `MultiQueryAttention` | Shared K/V across all heads |
+| `GroupedQueryAttention` | Groups of query heads sharing K/V (GQA) |
+| `MultiHeadLatentAttention` | Latent-space compression for keys/values |

transformer_lib-0.1.0/README.md

@@ -0,0 +1,33 @@
+# transformer-lib
+
+Attention mechanisms from Every_Thing_About_Transformer, as an importable Python library.
+
+## Installation
+
+```bash
+pip install -e .
+```
+
+## Usage
+
+```python
+from transformer_lib import MultiHeadAttention
+import torch
+
+mha = MultiHeadAttention(embed_dim=512, num_heads=8)
+x = torch.randn(2, 10, 512)
+out = mha(x, x, x)
+```
+
+## Available Modules
+
+| Module | Description |
+|---|---|
+| `SelfAttention` | Basic scaled dot-product self-attention |
+| `MultiHeadAttention` | Multi-head attention with separate Q/K/V projections |
+| `CausalAttention` | Multi-head with lower-triangular causal mask |
+| `CrossAttention` | Encoder-decoder cross-attention with self-attention pre-step |
+| `GlobalAttention` | Additive (Bahdanau-style) alignment scoring |
+| `MultiQueryAttention` | Shared K/V across all heads |
+| `GroupedQueryAttention` | Groups of query heads sharing K/V (GQA) |
+| `MultiHeadLatentAttention` | Latent-space compression for keys/values |

transformer_lib-0.1.0/examples/multi_head_attention_demo.py

@@ -0,0 +1,18 @@
+import torch
+from transformer_lib import MultiHeadAttention
+
+
+def demo_multi_head_attention(embed_dim=512, num_heads=8, batch_size=2, seq_len=10):
+    mha = MultiHeadAttention(embed_dim=embed_dim, num_heads=num_heads)
+
+    query = torch.randn(batch_size, seq_len, embed_dim)
+    key = torch.randn(batch_size, seq_len, embed_dim)
+    value = torch.randn(batch_size, seq_len, embed_dim)
+
+    attn_output = mha(query, key, value)
+    print(f"Output shape: {attn_output.shape}")
+    return attn_output
+
+
+if __name__ == "__main__":
+    demo_multi_head_attention()

transformer_lib-0.1.0/pyproject.toml

@@ -0,0 +1,22 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "transformer-lib"
+version = "0.1.0"
+description = "Attention mechanisms from Every_Thing_About_Transformer, as an importable library"
+readme = "README.md"
+requires-python = ">=3.9"
+license = {text = "MIT"}
+authors = [{name = "Aniket Verma"}]
+dependencies = [
+    "torch>=2.0",
+]
+
+[project.optional-dependencies]
+linear = ["linear-attention-transformer"]
+dev = ["pytest>=7.0"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/transformer_lib"]

transformer_lib-0.1.0/src/transformer_lib/__init__.py

@@ -0,0 +1,21 @@
+from .attention import (
+    SelfAttention,
+    MultiHeadAttention,
+    CausalAttention,
+    CrossAttention,
+    GlobalAttention,
+    MultiQueryAttention,
+    MultiHeadLatentAttention,
+    GroupedQueryAttention,
+)
+
+__all__ = [
+    "SelfAttention",
+    "MultiHeadAttention",
+    "CausalAttention",
+    "CrossAttention",
+    "GlobalAttention",
+    "MultiQueryAttention",
+    "MultiHeadLatentAttention",
+    "GroupedQueryAttention",
+]

transformer_lib-0.1.0/src/transformer_lib/attention/__init__.py

@@ -0,0 +1,19 @@
+from .self_attention import SelfAttention
+from .multi_head_attention import MultiHeadAttention
+from .causal_attention import CausalAttention
+from .cross_attention import CrossAttention
+from .global_attention import GlobalAttention
+from .multi_q_attention import MultiQueryAttention
+from .multi_head_latent_attention import MultiHeadLatentAttention
+from .grouped_q_attention import GroupedQueryAttention
+
+__all__ = [
+    "SelfAttention",
+    "MultiHeadAttention",
+    "CausalAttention",
+    "CrossAttention",
+    "GlobalAttention",
+    "MultiQueryAttention",
+    "MultiHeadLatentAttention",
+    "GroupedQueryAttention",
+]

transformer_lib-0.1.0/src/transformer_lib/attention/causal_attention.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CausalAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads, max_seq_len, dropout=0.1):
+        super().__init__()
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.d_k = embed_dim // num_heads
+        self.num_heads = num_heads
+
+        # Linear layers for query, key, and value
+        self.query = nn.Linear(embed_dim, embed_dim)
+        self.key = nn.Linear(embed_dim, embed_dim)
+        self.value = nn.Linear(embed_dim, embed_dim)
+        self.out = nn.Linear(embed_dim, embed_dim)
+
+        self.dropout = nn.Dropout(dropout)
+
+        # Register a buffer for the causal mask(lower triangular matrix)
+        # shape: (1, 1, max_seq_len, max_seq_len)
+        self.register_buffer("mask", torch.tril(torch.ones(max_seq_len, max_seq_len)).view(1, 1, max_seq_len, max_seq_len))
+
+    def forward(self, x):
+        batch_size, seq_len, embed_dim = x.size()
+
+        # Project inputs to query, key, and value
+        Q = self.query(x)  # (batch_size, seq_len, embed_dim)
+        K = self.key(x)    # (batch_size, seq_len, embed_dim)
+        V = self.value(x)  # (batch_size, seq_len, embed_dim)
+
+        # Reshape for multi-head attention
+        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)  # (batch_size, num_heads, seq_len, d_k)
+        K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)  # (batch_size, num_heads, seq_len, d_k)
+        V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)  # (batch_size, num_heads, seq_len, d_k)
+
+        # Scaled dot-product attention
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)  # (batch_size, num_heads, seq_len, seq_len)
+
+        # Apply causal mask
+        scores = scores.masked_fill(self.mask[:, :, :seq_len, :seq_len] == 0, float('-inf'))
+
+        # Softmax to get attention weights
+        attn_weights = F.softmax(scores, dim=-1)  # (batch_size, num_heads, seq_len, seq_len)
+        attn_weights = self.dropout(attn_weights)
+
+        # Compute the output
+        output = torch.matmul(attn_weights, V)  # (batch_size, num_heads, seq_len, d_k)
+
+        # Concatenate heads and project to output dimension
+        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)  # (batch_size, seq_len, embed_dim)
+        output = self.out(output)  # (batch_size, seq_len, embed_dim)
+
+        return output
+    
+
+"""The Causal Attention mechanism is a variant of multi-head attention that restricts a model 
+from attending to future tokens, ensuring strict autoregressive generation.  This is implemented by applying a
+ lower-triangular mask to the attention scores before the softmax function, effectively setting future positions to negative 
+ infinity."""

transformer_lib-0.1.0/src/transformer_lib/attention/cross_attention.py

@@ -0,0 +1,29 @@
+import torch
+import torch.nn as nn
+
+class CrossAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
+        self.cross_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.norm2 = nn.LayerNorm(embed_dim)
+
+    def forward(self, target, source, tgt_mask=None, src_mask=None):
+        # Self-attention on the target sequence
+        attn_output, _ = self.self_attn(target, target, target, attn_mask=tgt_mask)
+        target = self.norm1(target + attn_output)
+
+        # Cross-attention with the souce sequence
+        cross_output, _ = self.cross_attn(target, source, source, attn_mask=src_mask)
+        output = self.norm2(target + cross_output)
+
+        return output
+    
+
+
+"""Cross-attention
+
+Allows tokens from one sequence attend to a different sequence, bringing information from different sources or
+ modalities together (for example, description and image). It uses queries from one module – usually encoder, 
+ and keys and values from the other one (decoder)."""

transformer_lib-0.1.0/src/transformer_lib/attention/global_attention.py

@@ -0,0 +1,45 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class GlobalAttention(nn.Module):
+    def __init__(self, encoder_dim, decoder_dim):
+        super(GlobalAttention, self).__init__()
+        # Linear layers to project encoder and decoder states
+        self.encoder_linear = nn.Linear(encoder_dim, decoder_dim)
+        self.decoder_linear = nn.Linear(decoder_dim, decoder_dim)
+        self.attn_linear = nn.Linear(decoder_dim, 1)
+
+    def forward(self, encoder_outputs, decoder_hidden):
+        """
+        encoder_outputs: (batch_size, seq_len, encoder_dim)
+        decoder_hidden: (batch_size, decoder_dim)
+        """
+        # Project encoder outputs
+        # energy: (batch_size, seq_len, decoder_dim)
+        energy = self.encoder_linear(encoder_outputs)
+        
+        # Project decoder hidden state and expand to match sequence length
+        # decoder_hidden_expanded: (batch_size, 1, decoder_dim) -> (batch_size, seq_len, decoder_dim)
+        decoder_hidden_expanded = self.decoder_linear(decoder_hidden).unsqueeze(1)
+        
+        # Compute alignment scores
+        # attn_energies: (batch_size, seq_len, 1)
+        attn_energies = self.attn_linear(torch.tanh(energy + decoder_hidden_expanded))
+        
+        # Softmax to get attention weights
+        # attn_weights: (batch_size, seq_len, 1)
+        attn_weights = F.softmax(attn_energies, dim=1)
+        
+        # Compute context vector as weighted sum of encoder outputs
+        # context: (batch_size, encoder_dim)
+        context = torch.bmm(attn_weights.transpose(1, 2), encoder_outputs).squeeze(1)
+        
+        return context, attn_weights   
+    
+
+"""Global attention
+
+It is also a type of sparse attention that allows some tokens to attend to every token in a sequence,
+ while others have limited attention. Every token can also attend to these special tokens. 
+ It is often used with local attention to preserve the entire context handling. """

transformer_lib-0.1.0/src/transformer_lib/attention/grouped_q_attention.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+import math
+
+class GroupedQueryAttention(nn.Module):
+    def __init__(self, d_model, num_query_heads, num_kv_heads):
+        super().__init__()
+        assert d_model % num_query_heads == 0, "d_model must be divisible by num_query_heads"
+        assert num_query_heads % num_kv_heads == 0, "num_query_heads must be divisible by num_kv_heads"
+        
+        self.d_model = d_model
+        self.num_query_heads = num_query_heads
+        self.num_kv_heads = num_kv_heads
+        self.num_queries_per_kv = num_query_heads // num_kv_heads
+        self.head_dim = d_model // num_query_heads
+        
+        # Query projection: standard size
+        self.q_proj = nn.Linear(d_model, d_model)
+        
+        # Key/Value projections: reduced size (num_kv_heads * head_dim)
+        self.kv_dim = num_kv_heads * self.head_dim
+        self.k_proj = nn.Linear(d_model, self.kv_dim)
+        self.v_proj = nn.Linear(d_model, self.kv_dim)
+        
+        # Output projection
+        self.out_proj = nn.Linear(d_model, d_model)
+        
+        # Scaling factor
+        self.scale = math.sqrt(self.head_dim)
+
+    def forward(self, x, mask=None):
+        batch_size, seq_len, _ = x.shape
+        
+        # Projections
+        Q = self.q_proj(x)  # (batch, seq, d_model)
+        K = self.k_proj(x)  # (batch, seq, kv_dim)
+        V = self.v_proj(x)  # (batch, seq, kv_dim)
+        
+        # Reshape for multi-head attention
+        # Q: (batch, num_query_heads, seq, head_dim)
+        Q = Q.view(batch_size, seq_len, self.num_query_heads, self.head_dim).transpose(1, 2)
+        
+        # K/V: (batch, num_kv_heads, seq, head_dim)
+        K = K.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        V = V.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        
+        # Expand K and V to match number of query heads
+        # Each KV head is shared by num_queries_per_kv query heads
+        K = K.repeat_interleave(self.num_queries_per_kv, dim=1)
+        V = V.repeat_interleave(self.num_queries_per_kv, dim=1)
+        
+        # Attention scores
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
+        
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, float('-inf'))
+            
+        attn_weights = torch.softmax(scores, dim=-1)
+        attn_output = torch.matmul(attn_weights, V)
+        
+        # Concatenate heads
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, seq_len, self.d_model)
+        
+        return self.out_proj(attn_output)   

transformer_lib-0.1.0/src/transformer_lib/attention/linear_attention.py

@@ -0,0 +1,33 @@
+import torch
+from linear_attention_transformer import LinearAttentionTransformerLM
+
+
+def linear_attention_example(num_tokens=20000, dim=512, heads=8, max_seq_len=8192, device=None):
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = LinearAttentionTransformerLM(
+        num_tokens=num_tokens,
+        dim=dim,
+        heads=heads,
+        depth=1,
+        max_seq_len=max_seq_len,
+        causal=True,
+    ).to(device)
+
+    x = torch.randint(0, num_tokens, (1, max_seq_len), device=device)
+    output = model(x)
+    return output
+
+
+if __name__ == "__main__":
+    linear_attention_example()
+
+
+"""$ pip install linear-attention-transformer
+
+Linear attention is a faster version of self-attention that avoids comparing
+ all token pairs. By restructuring the computation with kernel functions 
+ (kernel functions are a way to compute similarity between vectors as if they were 
+ transformed into another space, without actually performing that transformation), 
+ it reduces complexity from O(N2) to O(N), making it much more efficient for long sequences."""

transformer_lib-0.1.0/src/transformer_lib/attention/multi_head_attention.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads):
+        super(MultiHeadAttention, self).__init__()
+        assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads"
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+
+        # Linear Layers for query, key, and value
+        self.query_linear = nn.Linear(embed_dim, embed_dim)
+        self.key_linear = nn.Linear(embed_dim, embed_dim)
+        self.value_linear = nn.Linear(embed_dim, embed_dim)
+
+        # output projection Layer
+        self.out_linear = nn.Linear(embed_dim, embed_dim)
+
+    def forward(self, query, key, value, mask=None):
+        batch_size = query.size(0)
+
+        # project and reshape to (batch_size, num_heads, seq_length, head_dim)
+        query = self.query_linear(query).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+        key = self.key_linear(key).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+        value = self.value_linear(value).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Scaled Dot-Product Attention
+        scores = torch.matmul(query, key.transpose(-2, -1)) / (self.head_dim ** 0.5)
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, float('-inf'))
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_output = torch.matmul(attn_weights, value)
+
+        # Concatenate heads and project
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.head_dim)
+        return self.out_linear(attn_output)
+    
+
+
+
+"""Multi-Head Attention (MHA)
+Splits attention into multiple heads (one independent attention mechanism), so the model can learn different interaction patterns
+ in parallel. It became the defining feature of Transformers and LLMs, and it still remains the baseline for novel attention 
+ variants"""

transformer_lib-0.1.0/src/transformer_lib/attention/multi_head_latent_attention.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class MultiHeadLatentAttention(nn.Module):
+    def __init__(self, d_model, num_heads, q_latent_dim, kv_latent_dim):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        
+        # Query compression and decomposition
+        self.w_q_d = nn.Linear(d_model, q_latent_dim, bias=False)
+        self.w_qk = nn.Linear(q_latent_dim, num_heads * kv_latent_dim, bias=False)
+        
+        # Key/Value compression and decomposition
+        self.w_kv_d = nn.Linear(d_model, kv_latent_dim, bias=False)
+        self.w_v_u = nn.Linear(kv_latent_dim, num_heads * self.head_dim, bias=False)
+        
+        # Output projection
+        self.w_o = nn.Linear(num_heads * self.head_dim, d_model)
+
+    def forward(self, x):
+        batch_size, seq_len, d_model = x.shape
+        
+        # 1. Project to latent space
+        c_q = self.w_q_d(x)  # [B, S, q_latent_dim]
+        c_kv = self.w_kv_d(x) # [B, S, kv_latent_dim]
+        
+        # 2. Compute attention scores via latent multiplication
+        # Reshape for multi-head: [B, S, num_heads, kv_latent_dim]
+        c_q_w_qk = self.w_qk(c_q).view(batch_size, seq_len, self.num_heads, -1)
+        
+        # Scores: [B, num_heads, S, S]
+        scores = torch.matmul(
+            c_q_w_qk.transpose(1, 2), 
+            c_kv.transpose(-2, -1)[:, None, ...]
+        ) / (self.head_dim ** 0.5)
+        
+        attn_weights = torch.softmax(scores, dim=-1)
+        
+        # 3. Restore Values and compute output
+        v = self.w_v_u(c_kv).view(batch_size, seq_len, self.num_heads, -1)
+        
+        # [B, num_heads, S, head_dim]
+        output = torch.matmul(attn_weights, v.transpose(1, 2)).transpose(1, 2).contiguous()
+        
+        # 4. Final projection
+        return self.w_o(output.view(batch_size, seq_len, -1))   

transformer_lib-0.1.0/src/transformer_lib/attention/multi_q_attention.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+import math
+
+class MultiQueryAttention(nn.Module):
+    def __init__(self, embedding_dim, num_heads):
+        super().__init__()
+        assert embedding_dim % num_heads == 0, "embedding_dim must be divisible by num_heads"
+        
+        self.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.head_dim = embedding_dim // num_heads
+        
+        # Each head has its own Query projection
+        self.W_q = nn.Linear(embedding_dim, embedding_dim, bias=False)
+        
+        # Single shared Key and Value projections for all heads
+        self.W_k = nn.Linear(embedding_dim, self.head_dim, bias=False)
+        self.W_v = nn.Linear(embedding_dim, self.head_dim, bias=False)
+        
+        # Final output projection
+        self.W_o = nn.Linear(embedding_dim, embedding_dim, bias=False)
+
+    def forward(self, x):
+        batch_size, seq_len, _ = x.size()
+        
+        # Compute Q, K, V
+        Q = self.W_q(x)  # [B, S, E]
+        K = self.W_k(x)  # [B, S, H]
+        V = self.W_v(x)  # [B, S, H]
+        
+        # Split Q into multiple heads: [B, S, E] -> [B, H, S, H]
+        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        
+        # Broadcast single K and V across all heads: [B, S, H] -> [B, 1, S, H]
+        K = K.unsqueeze(1)
+        V = V.unsqueeze(1)
+        
+        # Scaled Dot-Product Attention
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        attn_weights = torch.softmax(scores, dim=-1)
+        
+        # Apply attention weights to V
+        context = torch.matmul(attn_weights, V)  # [B, H, S, H]
+        
+        # Merge heads: [B, H, S, H] -> [B, S, E]
+        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embedding_dim)
+        
+        return self.W_o(context)   

transformer_lib-0.1.0/src/transformer_lib/attention/self_attention.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class SelfAttention(nn.Module):
+    def __init__(self, input_dim):
+        super().__init__()
+        self.input_dim = input_dim
+        self.query = nn.Linear(input_dim, input_dim)
+        self.key = nn.Linear(input_dim, input_dim)
+        self.value = nn.Linear(input_dim, input_dim)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        # x shape: (batch_size, seq_len, input_dim)
+        Q = self.query(x)
+        K = self.key(x)
+        V = self.value(x)
+
+        # Scaled dot-product attention
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.input_dim ** 0.5)
+
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, float('-inf'))
+
+        attention_weights = self.softmax(scores)
+        output = torch.matmul(attention_weights, V)
+        return output, attention_weights
+
+
+
+"""Self-attention:
+Let's each token attend directly to other tokens in same sequence: the model compares every token with
+the rest of the sequence and build a context-aware representation. This is what allows transformers to
+capture relationships accross distanceand make it possible to modellong-range dependencies without recurrence."""

transformer_lib-0.1.0/tests/test_attention.py

@@ -0,0 +1,101 @@
+import torch
+import pytest
+from transformer_lib import (
+    SelfAttention,
+    MultiHeadAttention,
+    CausalAttention,
+    CrossAttention,
+    GlobalAttention,
+    MultiQueryAttention,
+    MultiHeadLatentAttention,
+    GroupedQueryAttention,
+)
+
+
+class TestSelfAttention:
+    def test_output_shape(self):
+        model = SelfAttention(input_dim=64)
+        x = torch.randn(2, 10, 64)
+        output, weights = model(x)
+        assert output.shape == (2, 10, 64)
+        assert weights.shape == (2, 10, 10)
+
+    def test_with_mask(self):
+        model = SelfAttention(input_dim=64)
+        x = torch.randn(2, 10, 64)
+        mask = torch.ones(2, 10, 10)
+        output, weights = model(x, mask=mask)
+        assert output.shape == (2, 10, 64)
+
+
+class TestMultiHeadAttention:
+    def test_output_shape(self):
+        model = MultiHeadAttention(embed_dim=512, num_heads=8)
+        x = torch.randn(2, 10, 512)
+        out = model(x, x, x)
+        assert out.shape == (2, 10, 512)
+
+    def test_mask_applied(self):
+        model = MultiHeadAttention(embed_dim=64, num_heads=4)
+        x = torch.randn(2, 10, 64)
+        mask = torch.ones(2, 1, 10, 10)
+        out = model(x, x, x, mask=mask)
+        assert out.shape == (2, 10, 64)
+
+
+class TestCausalAttention:
+    def test_output_shape(self):
+        model = CausalAttention(embed_dim=512, num_heads=8, max_seq_len=100)
+        x = torch.randn(2, 10, 512)
+        out = model(x)
+        assert out.shape == (2, 10, 512)
+
+
+class TestCrossAttention:
+    def test_output_shape(self):
+        model = CrossAttention(embed_dim=512, num_heads=8)
+        target = torch.randn(2, 10, 512)
+        source = torch.randn(2, 20, 512)
+        out = model(target, source)
+        assert out.shape == (2, 10, 512)
+
+
+class TestGlobalAttention:
+    def test_output_shape(self):
+        model = GlobalAttention(encoder_dim=256, decoder_dim=512)
+        encoder_outputs = torch.randn(2, 15, 256)
+        decoder_hidden = torch.randn(2, 512)
+        context, weights = model(encoder_outputs, decoder_hidden)
+        assert context.shape == (2, 256)
+        assert weights.shape == (2, 15, 1)
+
+
+class TestMultiQueryAttention:
+    def test_output_shape(self):
+        model = MultiQueryAttention(embedding_dim=512, num_heads=8)
+        x = torch.randn(2, 10, 512)
+        out = model(x)
+        assert out.shape == (2, 10, 512)
+
+
+class TestMultiHeadLatentAttention:
+    def test_output_shape(self):
+        model = MultiHeadLatentAttention(d_model=512, num_heads=8, q_latent_dim=128, kv_latent_dim=64)
+        x = torch.randn(2, 10, 512)
+        out = model(x)
+        assert out.shape == (2, 10, 512)
+
+
+class TestGroupedQueryAttention:
+    def test_output_shape(self):
+        model = GroupedQueryAttention(d_model=512, num_query_heads=8, num_kv_heads=2)
+        x = torch.randn(2, 10, 512)
+        out = model(x)
+        assert out.shape == (2, 10, 512)
+
+    def test_with_mask(self):
+        model = GroupedQueryAttention(d_model=512, num_query_heads=8, num_kv_heads=2)
+        x = torch.randn(2, 10, 512)
+        mask = torch.ones(2, 1, 10, 10)
+        out = model(x, mask=mask)
+        assert out.shape == (2, 10, 512)