(
self,
x: mx.array,
memory: mx.array,
mask: mx.array,
memory_mask: mx.array,
cache: Optional[List[Tuple[mx.array, mx.array]]] = None,
)
| 247 | self.dense = DenseActivation(config) |
| 248 | |
| 249 | def __call__( |
| 250 | self, |
| 251 | x: mx.array, |
| 252 | memory: mx.array, |
| 253 | mask: mx.array, |
| 254 | memory_mask: mx.array, |
| 255 | cache: Optional[List[Tuple[mx.array, mx.array]]] = None, |
| 256 | ): |
| 257 | y = self.ln1(x) |
| 258 | y, cache = self.self_attention(y, y, y, mask, cache) |
| 259 | x = x + y |
| 260 | |
| 261 | y = self.ln2(x) |
| 262 | y, _ = self.cross_attention(y, memory, memory, memory_mask) |
| 263 | x = x + y |
| 264 | |
| 265 | y = self.ln3(x) |
| 266 | y = self.dense(y) |
| 267 | x = x + y |
| 268 | |
| 269 | return x, cache |
| 270 | |
| 271 | |
| 272 | class TransformerDecoder(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected