torch_blue.vi.VITransformerDecoderLayer

class torch_blue.vi.VITransformerDecoderLayer(d_model: int, nhead: int, dim_feedforward: int = 512, activation: torch.nn.Module = ReLU(), layer_norm_eps: float = 1e-05, norm_first: bool = False, batch_first: bool = True, bias: bool = True, variational_distribution: Distribution = MeanFieldNormal(), prior: Distribution = MeanFieldNormal(), rescale_prior: bool = True, kaiming_initialization: bool = True, prior_initialization: bool = False, return_log_probs: bool = True, device: torch.device | None = None, dtype: torch.dtype | None = None)

Bases: torch_blue.vi.base.VIModule

TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.

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

This does not support the dropout argument. In addition to all other arguments, this class accepts VIkwargs.

Parameters:

• torch_args – The same arguments and keyword arguments as the pytorch version TransformerDecoder (documentation here), except dropout, which is not relevant for BNNs.

• VIkwargs – Several standard keyword arguments. See VIkwargs for details.

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 = False, memory_is_causal: bool = False) → torch.Tensor

Pass the input through the decoder layer.

See documentation of nn.TranformerDecoderLayer for details.

This implementation also currently does not support the torch fastpath.