torch_blue.vi.VITransformer

class torch_blue.vi.VITransformer(d_model: int = 32, nhead: int = 1, num_encoder_layers: int = 1, num_decoder_layers: int = 1, dim_feedforward: int = 256, activation: torch.nn.Module = ReLU(), custom_encoder: Any | None = None, custom_decoder: Any | None = None, layer_norm_eps: float = 1e-05, batch_first: bool = False, 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

A Bayesian Transformer model.

Equivalent of nn.Transformer 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 Transformer (documentation here), except dropout, which is not relevant for BNNs.

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

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

Take in and process masked source/target sequences.

See documentation of nn.Transformer for details.

This implementation also currently does not support the torch fastpath.