Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add is_causal API for TransformerDecoder #97166

Closed
wants to merge 2 commits into from
Closed

Add is_causal API for TransformerDecoder #97166

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
547b851
Add `is_causal` API for `TransformerDecoder`
janEbert Mar 20, 2023
8808120
[noci] Fix wrong word
janEbert Mar 24, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
13 changes: 10 additions & 3 deletions torch/nn/modules/transformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,8 @@ def __init__(self, decoder_layer, num_layers, norm=None):

def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,
memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
memory_key_padding_mask: Optional[Tensor] = None, tgt_is_causal: bool = False,
memory_is_causal: bool = False) -> Tensor:
r"""Pass the inputs (and mask) through the decoder layer in turn.

Args:
Expand All @@ -359,6 +360,10 @@ def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None
memory_mask: the mask for the memory sequence (optional).
tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
memory_key_padding_mask: the mask for the memory keys per batch (optional).
tgt_is_causal: If specified, applies a causal mask as tgt mask.
Mutually exclusive with providing tgt_mask. Default: ``False``.
memory_is_causal: If specified, applies a causal mask as memory mask.
Mutually exclusive with providing memory_mask. Default: ``False``.

Shape:
see the docs in Transformer class.
Expand All @@ -369,7 +374,9 @@ def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None
output = mod(output, memory, tgt_mask=tgt_mask,
memory_mask=memory_mask,
tgt_key_padding_mask=tgt_key_padding_mask,
memory_key_padding_mask=memory_key_padding_mask)
memory_key_padding_mask=memory_key_padding_mask,
tgt_is_causal=tgt_is_causal,
memory_is_causal=memory_is_causal)

if self.norm is not None:
output = self.norm(output)
Expand Down Expand Up @@ -700,7 +707,7 @@ def forward(
memory_key_padding_mask: the mask for the memory keys per batch (optional).
tgt_is_causal: If specified, applies a causal mask as tgt mask.
Mutually exclusive with providing tgt_mask. Default: ``False``.
memory_is_causal: If specified, applies a causal mask as tgt mask.
memory_is_causal: If specified, applies a causal mask as memory mask.
Mutually exclusive with providing memory_mask. Default: ``False``.
Shape:
see the docs in Transformer class.
Expand Down