torch_blue.vi.VITransformerEncoderLayer

class torch_blue.vi.VITransformerEncoderLayer(d_model: int, nhead: int, dim_feedforward: int = 512, activation: torch.nn.Module = ReLU(), layer_norm_eps: float = 1e-05, batch_first: bool = True, norm_first: bool = False, 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

TransformerEncoderLayer is made up of self-attn and feedforward network.

Equivalent of nn.TransformerEncoderLayer 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 TransformerEncoderLayer (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(src: torch.Tensor, src_mask: torch.Tensor | None = None, src_key_padding_mask: torch.Tensor | None = None, is_causal: bool = False) → torch.Tensor

Pass the input through the encoder layer.

See documentation of nn.TranformerEncoderLayer for details.

This implementation also currently does not support the torch fastpath.