torch_blue.vi.VIkwargs

class torch_blue.vi.VIkwargs

Bases: TypedDict

Common keyword arguments for most VIModules.

Most VIModule accept these as keyword arguments. Unless specified otherwise, they use the defaults below.

This class is only used for documentation and type hinting.

Parameters:

  • variational_distribution (Union[Distribution, List[Distribution]], default: MeanFieldNormal()) – Either one torch_blue.vi.distributions.Distribution , which is used for all random variables, or a list of them, one for each random variable. This specifies the assumed parametrization of the weight distribution.

  • prior (Union[Distribution, List[Distribution]], default: MeanFieldNormal()) – Either one Distribution , which is used for all random variables, or a list of them, one for each random variable. This specifies the previous knowledge about the weight distribution.

  • rescale_prior (bool, default: False) – If True , the priors _scaling_parameters are scaled with the sqrt of the layer width. This may be necessary to maintain normalization for wide layers.

  • prior_initialization (bool, default: False) – If True parameters are initialized according to the prior. If False parameters are initialized similar to non-Bayesian networks.

  • return_log_probs (bool, default: True) – If True the model forward pass returns the log probability of the sampled weights. This is required for use of KullbackLeiblerLoss.

  • device (Optional[torch.device], default: None) – The torch.device on which the module should be stored.

  • dtype (Optional[torch.dtype], default: None) – The torch.dtype of the module parameters.