| 606 | c, 4 * c, padding='LEFT', kernel_size=kernel_size, dropout=relu_dropout, act=act) |
| 607 | |
| 608 | def forward( |
| 609 | self, |
| 610 | x, |
| 611 | encoder_out=None, |
| 612 | encoder_padding_mask=None, |
| 613 | incremental_state=None, |
| 614 | self_attn_mask=None, |
| 615 | self_attn_padding_mask=None, |
| 616 | attn_out=None, |
| 617 | reset_attn_weight=None, |
| 618 | **kwargs, |
| 619 | ): |
| 620 | layer_norm_training = kwargs.get('layer_norm_training', None) |
| 621 | if layer_norm_training is not None: |
| 622 | self.layer_norm1.training = layer_norm_training |
| 623 | self.layer_norm2.training = layer_norm_training |
| 624 | self.layer_norm3.training = layer_norm_training |
| 625 | residual = x |
| 626 | x = self.layer_norm1(x) |
| 627 | x, _ = self.self_attn( |
| 628 | query=x, |
| 629 | key=x, |
| 630 | value=x, |
| 631 | key_padding_mask=self_attn_padding_mask, |
| 632 | incremental_state=incremental_state, |
| 633 | attn_mask=self_attn_mask |
| 634 | ) |
| 635 | x = F.dropout(x, self.dropout, training=self.training) |
| 636 | x = residual + x |
| 637 | |
| 638 | residual = x |
| 639 | x = self.layer_norm2(x) |
| 640 | if encoder_out is not None: |
| 641 | x, attn = self.encoder_attn( |
| 642 | query=x, |
| 643 | key=encoder_out, |
| 644 | value=encoder_out, |
| 645 | key_padding_mask=encoder_padding_mask, |
| 646 | incremental_state=incremental_state, |
| 647 | static_kv=True, |
| 648 | enc_dec_attn_constraint_mask=None, #utils.get_incremental_state(self, incremental_state, 'enc_dec_attn_constraint_mask'), |
| 649 | reset_attn_weight=reset_attn_weight |
| 650 | ) |
| 651 | attn_logits = attn[1] |
| 652 | else: |
| 653 | assert attn_out is not None |
| 654 | x = self.encoder_attn.in_proj_v(attn_out.transpose(0, 1)) |
| 655 | attn_logits = None |
| 656 | x = F.dropout(x, self.dropout, training=self.training) |
| 657 | x = residual + x |
| 658 | |
| 659 | residual = x |
| 660 | x = self.layer_norm3(x) |
| 661 | x = self.ffn(x, incremental_state=incremental_state) |
| 662 | x = F.dropout(x, self.dropout, training=self.training) |
| 663 | x = residual + x |
| 664 | # if len(attn_logits.size()) > 3: |
| 665 | # indices = attn_logits.softmax(-1).max(-1).values.sum(-1).argmax(-1) |