nessai.flows.maf

Implementation of MaskedAutoregressiveFlow.

Module Contents

Classes

MaskedAutoregressiveFlow

Autoregressive flow with masked coupling transforms.

Attributes

logger

nessai.flows.maf.logger

class nessai.flows.maf.MaskedAutoregressiveFlow(features, hidden_features, num_layers, num_blocks_per_layer, context_features=None, use_residual_blocks=True, use_random_masks=False, use_random_permutations=False, activation=F.relu, dropout_probability=0.0, batch_norm_within_layers=False, batch_norm_between_layers=False)

Bases: nessai.flows.base.NFlow

Autoregressive flow with masked coupling transforms.

Based on the implementation from nflows: https://github.com/bayesiains/nflows/blob/master/nflows/flows/autoregressive.py

but also included context features.

Parameters

featuresint

Number of features (dimensions) in the data space

hidden_featuresint

Number of neurons per layer in each neural network

num_layersint

Number of coupling transformations

num_blocks_per_layerint

Number of layers (or blocks for resnet) per neural network for each coupling transform

context_featuresint, optional

Number of context (conditional) parameters.

use_residual_blocksbool, optional

Use residual blocks in the MADE network.

use_random_masksbool, optional

Use random masks in the MADE network.

use_random_permutationbool, optional

Use a random permutation instead of the default reverse permutation.

activationfunction, optional

Activation function implemented in torch.

dropout_probabilityfloat, optional

Dropout probability used in each layer of the neural network

batch_norm_within_layersbool, optional

Enable or disable batch norm within the neural network for each coupling transform

batch_norm_between_layersbool, optional

Enable or disable batch norm between coupling transforms