| 198 | |
| 199 | |
| 200 | class TransformerDecoderLayer(nn.Module): |
| 201 | def __init__(self, config: T5Config): |
| 202 | super().__init__() |
| 203 | self.self_attention = MultiHeadAttention(config) |
| 204 | self.cross_attention = MultiHeadAttention(config) |
| 205 | self.ln1 = nn.RMSNorm(config.d_model, eps=config.layer_norm_epsilon) |
| 206 | self.ln2 = nn.RMSNorm(config.d_model, eps=config.layer_norm_epsilon) |
| 207 | self.ln3 = nn.RMSNorm(config.d_model, eps=config.layer_norm_epsilon) |
| 208 | self.dense = DenseActivation(config) |
| 209 | |
| 210 | def __call__( |
| 211 | self, |
| 212 | x: mx.array, |
| 213 | memory: mx.array, |
| 214 | mask: mx.array, |
| 215 | memory_mask: mx.array, |
| 216 | cache: Optional[Tuple[mx.array, mx.array]] = None, |
| 217 | ) -> Tuple[mx.array, Tuple[mx.array, mx.array]]: |
| 218 | y = self.ln1(x) |
| 219 | y, new_cache = self.self_attention(y, y, y, mask, cache) |
| 220 | x = x + y |
| 221 | |
| 222 | y = self.ln2(x) |
| 223 | y, _ = self.cross_attention(y, memory, memory, memory_mask) |
| 224 | x = x + y |
| 225 | |
| 226 | y = self.ln3(x) |
| 227 | y = self.dense(y) |
| 228 | x = x + y |
| 229 | |
| 230 | return x, new_cache |
| 231 | |
| 232 | |
| 233 | def create_additive_causal_mask(N: int, offset: int = 0): |