torch_blue.vi.VILinear

class torch_blue.vi.VILinear(in_features: int, out_features: int, variational_distribution: Distribution | List[Distribution] = MeanFieldNormal(), prior: Distribution | List[Distribution] = MeanFieldNormal(), bias: bool = True, rescale_prior: bool = False, 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

Applies an affine linear transformation to the incoming data: \(y = xA^T + b\).

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

In addition to those arguments, this class accepts VIkwargs.

This module’s random variables are

  • (“weight”, “bias”) if bias == True

  • (“weight”, ) if bias == False

Parameters:

  • torch_args – The same arguments and keyword arguments as the pytorch version Linear (documentation here)

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

forward(input_: torch.Tensor) → torch.Tensor

Forward computation.

Parameters:

input (Tensor) – Input tensor of shape (*, in_features).

Returns:

output – Output tensor of shape (*, out_features). Auto-sampling will add a sample dimension at the start for the overall output.

Return type:

Tensor