| 33 | |
| 34 | |
| 35 | class TransformerBlock(nn.Module): |
| 36 | def __init__( |
| 37 | self, |
| 38 | model_dims: int, |
| 39 | num_heads: int, |
| 40 | hidden_dims: Optional[int] = None, |
| 41 | memory_dims: Optional[int] = None, |
| 42 | ): |
| 43 | super().__init__() |
| 44 | |
| 45 | self.norm1 = nn.LayerNorm(model_dims) |
| 46 | self.attn1 = nn.MultiHeadAttention(model_dims, num_heads) |
| 47 | self.attn1.out_proj.bias = mx.zeros(model_dims) |
| 48 | |
| 49 | memory_dims = memory_dims or model_dims |
| 50 | self.norm2 = nn.LayerNorm(model_dims) |
| 51 | self.attn2 = nn.MultiHeadAttention( |
| 52 | model_dims, num_heads, key_input_dims=memory_dims |
| 53 | ) |
| 54 | self.attn2.out_proj.bias = mx.zeros(model_dims) |
| 55 | |
| 56 | hidden_dims = hidden_dims or 4 * model_dims |
| 57 | self.norm3 = nn.LayerNorm(model_dims) |
| 58 | self.linear1 = nn.Linear(model_dims, hidden_dims) |
| 59 | self.linear2 = nn.Linear(model_dims, hidden_dims) |
| 60 | self.linear3 = nn.Linear(hidden_dims, model_dims) |
| 61 | |
| 62 | def __call__(self, x, memory, attn_mask, memory_mask): |
| 63 | # Self attention |
| 64 | y = self.norm1(x) |
| 65 | y = self.attn1(y, y, y, attn_mask) |
| 66 | x = x + y |
| 67 | |
| 68 | # Cross attention |
| 69 | y = self.norm2(x) |
| 70 | y = self.attn2(y, memory, memory, memory_mask) |
| 71 | x = x + y |
| 72 | |
| 73 | # FFN |
| 74 | y = self.norm3(x) |
| 75 | y_a = self.linear1(y) |
| 76 | y_b = self.linear2(y) |
| 77 | y = y_a * nn.gelu(y_b) |
| 78 | y = self.linear3(y) |
| 79 | x = x + y |
| 80 | |
| 81 | return x |
| 82 | |
| 83 | |
| 84 | class Transformer2D(nn.Module): |