(self, hidden_states, encoder_states, ltor_mask, cross_mask=None)
| 465 | output_layer_init_method=output_layer_init_method) |
| 466 | |
| 467 | def forward(self, hidden_states, encoder_states, ltor_mask, cross_mask=None): |
| 468 | # hidden_states: [b, s, h] |
| 469 | # ltor_mask: [1, 1, s, s] |
| 470 | |
| 471 | # Layer norm at the begining of the transformer layer. |
| 472 | layernorm_output = self.input_layernorm(hidden_states) |
| 473 | # Self attention. |
| 474 | self_attention_output = self.self_attention(layernorm_output, ltor_mask) |
| 475 | # Residual connection. |
| 476 | self_layernorm_input = hidden_states + self_attention_output |
| 477 | # Layer norm post the self attention. |
| 478 | self_layernorm_output = self.post_self_layernorm(self_layernorm_input) |
| 479 | # Cross attention |
| 480 | attention_output = self.cross_attention(self_layernorm_output, encoder_states, cross_mask) |
| 481 | # Residual connection |
| 482 | layernorm_input = self_layernorm_input + attention_output |
| 483 | # Layer norm post the cross attention |
| 484 | layernorm_output = self.post_attention_layernorm(layernorm_input) |
| 485 | # MLP. |
| 486 | mlp_output = self.mlp(layernorm_output) |
| 487 | # Second residual connection. |
| 488 | output = layernorm_input + mlp_output |
| 489 | return output |
| 490 | |
| 491 | |
| 492 | class ParallelTransformerLayer(torch.nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected