Automatic Differentiation Variational Inference (ADVI)#

class fortuna.prob_model.posterior.normalizing_flow.advi.advi_approximator.ADVIPosteriorApproximator(std_init_params=0.1, log_std_base=-2.3, n_loss_samples=3)[source]#

Automatic Differentiation Variational Inference (ADVI) approximator. It is responsible to define how the posterior distribution is approximated.

Parameters:

  • std_init_params (float) – The standard deviation of the Gaussian distribution used to initialize the parameters of the flow.

  • log_std_base (float) – The normalizing flow transforms a base distribution into an approximation of the posterior. The base distribution is assumed to be an isotropic Gaussian, with this argument as the log-standard deviation.

  • n_loss_samples (int) – Number of samples to approximate the loss, that is the KL divergence (or the ELBO, equivalently).

property posterior_method_kwargs: Dict[str, Any]#
class fortuna.prob_model.posterior.normalizing_flow.advi.advi_posterior.ADVIPosterior(joint, posterior_approximator)[source]#

Bases: Posterior

Automatic Differentiation Variational Inference (ADVI) approximate posterior class.

Parameters:

  • joint (Joint) – A joint distribution object.

  • posterior_approximator (ADVI) – An ADVI posterior approximator.

fit(train_data_loader, val_data_loader=None, fit_config=FitConfig(), map_fit_config=None, **kwargs)[source]#

Fit the posterior distribution. A posterior state will be internally stored.

Parameters:

  • train_data_loader (DataLoader) – Training data loader.

  • val_data_loader (Optional[DataLoader]) – Validation data loader.

  • fit_config (FitConfig) – A configuration object.

Returns:

A status including metrics describing the fitting process.

Return type:

Status

load_state(checkpoint_path)[source]#

Load the state of the posterior distribution from a checkpoint path. The checkpoint must be compatible with the current probabilistic model.

Parameters:

checkpoint_path (Path) – Path to checkpoint file or directory to restore.

Return type:

None

sample(rng=None, input_shape=None, inputs_loader=None, inputs=None, **kwargs)[source]#

Sample from the posterior distribution. Either input_shape or _inputs_loader must be passed. :type rng: Optional[PRNGKeyArray] :param rng: A random number generator. If not passed, this will be taken from the attributes of this class. :type rng: Optional[PRNGKeyArray] :type input_shape: Optional[Tuple[int, ...]] :param input_shape: Shape of a single input. :type input_shape: Optional[Tuple[int, …]] :type inputs_loader: Optional[InputsLoader] :param inputs_loader: Input data loader. If input_shape is passed, then inputs and inputs_loader are ignored. :type inputs_loader: Optional[InputsLoader] :type inputs: Union[Array, ndarray, None] :param inputs: Input variables. :type inputs: Optional[Array]

Returns:

A sample from the posterior distribution.

Return type:

JointState

class fortuna.prob_model.posterior.normalizing_flow.advi.advi_state.ADVIState(step, apply_fn, params, tx, opt_state, encoded_name=(65, 68, 86, 73, 83, 116, 97, 116, 101), frozen_params=None, dynamic_scale=None, mutable=None, calib_params=None, calib_mutable=None, grad_accumulated=None, _encoded_which_params=None)[source]#

Bases: NormalizingFlowState

encoded_name#

ADVI state name encoded as an array.

Type:

jnp.ndarray

classmethod init(params, mutable=None, optimizer=None, calib_params=None, calib_mutable=None, grad_accumulated=None, dynamic_scale=None, **kwargs)#

Initialize a posterior distribution state.

Parameters:

  • params (Params) – The parameters characterizing an approximation of the posterior distribution.

  • optimizer (Optional[OptaxOptimizer]) – An Optax optimizer associated with the posterior state.

  • mutable (Optional[Mutable]) – The mutable objects characterizing an approximation of the posterior distribution.

  • calib_params (Optional[CalibParams]) – The parameters objects characterizing an approximation of the posterior distribution.

  • calib_mutable (Optional[CalibMutable]) – The calibration mutable objects characterizing an approximation of the posterior distribution.

  • grad_accumulated (Optional[jnp.ndarray]) – The gradients accumulated in consecutive training steps (used only when gradient_accumulation_steps > 1).

  • dynamic_scale (Optional[dynamic_scale.DynamicScale]) – Dynamic loss scaling for mixed precision gradients.

Returns:

A posterior distribution state.

Return type:

Any

classmethod init_from_dict(d, optimizer=None, **kwargs)#

Initialize a posterior distribution state from a dictionary.

Parameters:

  • d (Dict) – A dictionary including attributes of the posterior state.

  • optimizer (Optional[OptaxOptimizer]) – An optax optimizer to assign to the posterior state.

Returns:

A posterior state.

Return type:

PosteriorState