(self,
hidden_size,
num_attention_heads,
attention_dropout_prob,
output_dropout_prob,
layernorm_epsilon,
init_method,
output_layer_init_method=None)
| 418 | """ |
| 419 | |
| 420 | def __init__(self, |
| 421 | hidden_size, |
| 422 | num_attention_heads, |
| 423 | attention_dropout_prob, |
| 424 | output_dropout_prob, |
| 425 | layernorm_epsilon, |
| 426 | init_method, |
| 427 | output_layer_init_method=None): |
| 428 | super(ParallelDecoderLayer, self).__init__() |
| 429 | # Set output layer initialization if not provided. |
| 430 | if output_layer_init_method is None: |
| 431 | output_layer_init_method = init_method |
| 432 | |
| 433 | # Layernorm on the input data. |
| 434 | self.input_layernorm = LayerNorm(hidden_size, eps=layernorm_epsilon) |
| 435 | |
| 436 | # Self attention. |
| 437 | self.self_attention = ParallelSelfAttention( |
| 438 | hidden_size, |
| 439 | num_attention_heads, |
| 440 | attention_dropout_prob, |
| 441 | output_dropout_prob, |
| 442 | init_method, |
| 443 | output_layer_init_method=output_layer_init_method) |
| 444 | |
| 445 | # Layernorm after the self attention. |
| 446 | self.post_self_layernorm = LayerNorm(hidden_size, eps=layernorm_epsilon) |
| 447 | |
| 448 | self.cross_attention = ParallelCrossAttention( |
| 449 | hidden_size, |
| 450 | num_attention_heads, |
| 451 | attention_dropout_prob, |
| 452 | output_dropout_prob, |
| 453 | init_method, |
| 454 | output_layer_init_method=output_layer_init_method |
| 455 | ) |
| 456 | |
| 457 | # Layernorm after the cross attention. |
| 458 | self.post_attention_layernorm = LayerNorm(hidden_size, eps=layernorm_epsilon) |
| 459 | |
| 460 | # MLP |
| 461 | self.mlp = ParallelMLP( |
| 462 | hidden_size, |
| 463 | output_dropout_prob, |
| 464 | init_method, |
| 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] |
nothing calls this directly
no test coverage detected