bayesflow.coupling_networks module

class bayesflow.coupling_networks.AffineCoupling(*args, **kwargs)

Bases: Model

Implements a conditional affine coupling block according to [1, 2], with additional options, such as residual blocks or Monte Carlo Dropout.

[1] Kingma, D. P., & Dhariwal, P. (2018). Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems, 31.

[2] Ardizzone, L., Lüth, C., Kruse, J., Rother, C., & Köthe, U. (2019). Guided image generation with conditional invertible neural networks. arXiv preprint arXiv:1907.02392.

__init__(dim_out, settings_dict, **kwargs)

Creates one half of an affine coupling layer to be used as part of a CouplingLayer in an InvertibleNetwork instance.

Parameters:

dim_outint

The output dimensionality of the affine coupling layer.

settings_dictdict

The settings for the inner networks. Defaults will use: ``settings_dict={

“dense_args” : dict(units=128, activation=”relu”), “num_dense” : 2, “spec_norm” : False, “mc_dropout” : False, “dropout” : True, “residual” : False, “dropout_prob” : 0.01, “soft_clamping” : 1.9

call(split1, split2, condition, inverse=False, **kwargs)

Performs one pass through an affine coupling layer (either inverse or forward).

Parameters:

split1tf.Tensor of shape (batch_size, …, input_dim//2)

The first partition of the input vector(s)

split2tf.Tensor of shape (batch_size, …, ceil[input_dim//2])

The second partition of the input vector(s)

conditiontf.Tensor or None

The conditioning data of interest, for instance, x = summary_fun(x), shape (batch_size, …). If condition is None, then the layer reduces to an unconditional coupling.

inversebool, optional, default: False

Flag indicating whether to run the block forward or backward.

Returns:

(z, log_det_J)tuple(tf.Tensor, tf.Tensor)

If inverse=False: The transformed input and the corresponding Jacobian of the transformation, z shape: (batch_size, …, input_dim//2), log_det_J shape: (batch_size, …)

targettf.Tensor

If inverse=True: The back-transformed z, shape (batch_size, …, inp_dim//2)

class bayesflow.coupling_networks.SplineCoupling(*args, **kwargs)

Bases: Model

Implements a conditional spline coupling block according to [1, 2], with additional options, such as residual blocks or Monte Carlo Dropout.

[1] Durkan, C., Bekasov, A., Murray, I., & Papamakarios, G. (2019). Neural spline flows. Advances in Neural Information Processing Systems, 32.

[2] Ardizzone, L., Lüth, C., Kruse, J., Rother, C., & Köthe, U. (2019). Guided image generation with conditional invertible neural networks. arXiv preprint arXiv:1907.02392.

Implement only rational quadratic splines (RQS), since these appear to work best in practice and lead to stable training.

__init__(dim_out, settings_dict, **kwargs)

Creates one half of a spline coupling layer to be used as part of a CouplingLayer in an InvertibleNetwork instance.

Parameters:

dim_outint

The output dimensionality of the coupling layer.

settings_dictdict

The settings for the inner networks. Defaults will use: ``settings_dict={

“dense_args” : dict(units=128, activation=”relu”), “num_dense” : 2, “spec_norm” : False, “mc_dropout” : False, “dropout” : True, “residual” : False, “dropout_prob” : 0.05, “bins” : 16, “default_domain” : (-5., 5., -5., 5.)

call(split1, split2, condition, inverse=False, **kwargs)

Performs one pass through a spline coupling layer (either inverse or forward).

Parameters:

split1tf.Tensor of shape (batch_size, …, input_dim//2)

The first partition of the input vector(s)

split2tf.Tensor of shape (batch_size, …, ceil[input_dim//2])

The second partition of the input vector(s)

conditiontf.Tensor or None

The conditioning data of interest, for instance, x = summary_fun(x), shape (batch_size, …). If condition is None, then the layer recuces to an unconditional coupling.

inversebool, optional, default: False

Flag indicating whether to run the block forward or backward.

Returns:

(z, log_det_J)tuple(tf.Tensor, tf.Tensor)

If inverse=False: The transformed input and the corresponding Jacobian of the transformation, z shape: (batch_size, …, input_dim//2), log_det_J shape: (batch_size, …)

targettf.Tensor

If inverse=True: The back-transformed z, shape (batch_size, …, inp_dim//2)

class bayesflow.coupling_networks.CouplingLayer(*args, **kwargs)

Bases: Model

General wrapper for a coupling layer (either affine or spline) with different settings.

__init__(latent_dim, coupling_settings=None, coupling_design='affine', permutation='fixed', use_act_norm=True, act_norm_init=None, **kwargs)

Creates an invertible coupling layers instance with the provided hyperparameters.

Parameters:

latent_dimint

The dimensionality of the latent space (equal to the dimensionality of the target variable)

coupling_settingsdict or None, optional, default: None

The coupling network settings to pass to the internal coupling layers. See default_settings for the required entries.

coupling_designstr or callable, optional, default: ‘affine’

The type of internal coupling network to use. Must be in [‘affine’, ‘spline’]. In general, spline couplings run slower than affine couplings, but requires fewer coupling layers. Spline couplings may work best with complex (e.g., multimodal) low-dimensional problems. The difference will become less and less pronounced as we move to higher dimensions.

permutationstr or None, optional, default: ‘fixed’

Whether to use permutations between coupling layers. Highly recommended if num_coupling_layers > 1 Important: Must be in [‘fixed’, ‘learnable’, None]

use_act_normbool, optional, default: True

Whether to use activation normalization after each coupling layer. Recommended to keep default.

act_norm_initnp.ndarray of shape (num_simulations, num_params) or None, optional, default: None

Optional data-dependent initialization for the internal ActNorm layers.

**kwargsdict

Optional keyword arguments (e.g., name) passed to the tf.keras.Model __init__ method.

call(target_or_z, condition, inverse=False, **kwargs)

Performs one pass through the affine coupling layer (either inverse or forward).

Parameters:

target_or_ztf.Tensor

The estimation quantities of interest or latent representations z ~ p(z), shape (batch_size, …)

conditiontf.Tensor or None

The conditioning data of interest, for instance, x = summary_fun(x), shape (batch_size, …). If condition is None, then the layer reduces to an unconditional ACL.

inversebool, optional, default: False

Flag indicating whether to run the block forward or backward.

Returns:

(z, log_det_J)tuple(tf.Tensor, tf.Tensor)

If inverse=False: The transformed input and the corresponding Jacobian of the transformation, z shape: (batch_size, inp_dim), log_det_J shape: (batch_size, )

targettf.Tensor

If inverse=True: The back-transformed z, shape (batch_size, inp_dim)

Notes

If inverse=False, the return is (z, log_det_J).

If inverse=True, the return is target

forward(target, condition, **kwargs)

Performs a forward pass through a coupling layer with an optinal Permutation and ActNorm layers.

Parameters:

targettf.Tensor

The estimation quantities of interest, for instance, parameter vector of shape (batch_size, theta_dim)

conditiontf.Tensor or None

The conditioning vector of interest, for instance, x = summary(x), shape (batch_size, summary_dim) If None, transformation amounts to unconditional estimation.

Returns:

(z, log_det_J)tuple(tf.Tensor, tf.Tensor)

The transformed input and the corresponding Jacobian of the transformation.

inverse(latent, condition, **kwargs)

Performs an inverse pass through a coupling layer with an optinal Permutation and ActNorm layers.

Parameters:

ztf.Tensor

latent variables z ~ p(z), shape (batch_size, theta_dim)

conditiontf.Tensor or None

The conditioning vector of interest, for instance, x = summary(x), shape (batch_size, summary_dim). If None, transformation amounts to unconditional estimation.

Returns:

targettf.Tensor

The back-transformed latent variable z.