torch_blue.vi.VITransformerDecoder

class torch_blue.vi.VITransformerDecoder(decoder_layer: VITransformerDecoderLayer, num_layers: int, norm: torch.nn.Module | None = None, return_log_probs: bool = True)

Bases: torch_blue.vi.base.VIModule

TransformerDecoder is a stack of N decoder layers.

Equivalent of nn.TransformerDecoder with variational inference. See its documentation for usage.

return_log_probs = True

Set whether the module returns log probabilities.

Log probabilities are required for most standard losses.

Returns:

Whether the module returns log probabilities.

Return type:

bool

forward(tgt: torch.Tensor, memory: torch.Tensor, tgt_mask: torch.Tensor | None = None, memory_mask: torch.Tensor | None = None, tgt_key_padding_mask: torch.Tensor | None = None, memory_key_padding_mask: torch.Tensor | None = None, tgt_is_causal: bool | None = None, memory_is_causal: bool = False) → torch.Tensor

Pass the input through the decoder layers in turn.

See documentation of nn.TransformerDecoder for details.

This implementation also currently does not support the torch fastpath.