(self,
hidden_size,
num_attention_heads,
attention_dropout_prob,
output_dropout_prob,
layernorm_epsilon,
init_method,
output_layer_init_method=None,
relative_encoding=False,
performer=False,
attention_scale=1.0)
| 519 | """ |
| 520 | |
| 521 | def __init__(self, |
| 522 | hidden_size, |
| 523 | num_attention_heads, |
| 524 | attention_dropout_prob, |
| 525 | output_dropout_prob, |
| 526 | layernorm_epsilon, |
| 527 | init_method, |
| 528 | output_layer_init_method=None, |
| 529 | relative_encoding=False, |
| 530 | performer=False, |
| 531 | attention_scale=1.0): |
| 532 | super(ParallelTransformerLayer, self).__init__() |
| 533 | # Set output layer initialization if not provided. |
| 534 | if output_layer_init_method is None: |
| 535 | output_layer_init_method = init_method |
| 536 | |
| 537 | # Layernorm on the input data. |
| 538 | self.input_layernorm = LayerNorm(hidden_size, eps=layernorm_epsilon) |
| 539 | |
| 540 | # Self attention. |
| 541 | self.attention = ParallelSelfAttention( |
| 542 | hidden_size, |
| 543 | num_attention_heads, |
| 544 | attention_dropout_prob, |
| 545 | output_dropout_prob, |
| 546 | init_method, |
| 547 | output_layer_init_method=output_layer_init_method, |
| 548 | relative_encoding=relative_encoding, |
| 549 | performer=performer, |
| 550 | attention_scale=attention_scale) |
| 551 | |
| 552 | # Layernorm on the input data. |
| 553 | self.post_attention_layernorm = LayerNorm(hidden_size, |
| 554 | eps=layernorm_epsilon) |
| 555 | |
| 556 | # MLP |
| 557 | self.mlp = ParallelMLP( |
| 558 | hidden_size, |
| 559 | output_dropout_prob, |
| 560 | init_method, |
| 561 | output_layer_init_method=output_layer_init_method) |
| 562 | |
| 563 | def forward(self, hidden_states, ltor_mask, position_embeddings=None, r_w_bias=None, r_r_bias=None, mem=None): |
| 564 | # hidden_states: [b, s, h] |
nothing calls this directly
no test coverage detected