core.nn.nsf.create_base_transform

core.nn.nsf.create_base_transform(
    i,
    param_dim,
    context_dim=None,
    hidden_dim=512,
    num_transform_blocks=2,
    activation='relu',
    dropout_probability=0.0,
    batch_norm=False,
    num_bins=8,
    tail_bound=1.0,
    apply_unconditional_transform=False,
    base_transform_type='rq-coupling',
)

Build a base NSF transform of y, conditioned on x.

This uses the PiecewiseRationalQuadraticCoupling transform or the MaskedPiecewiseRationalQuadraticAutoregressiveTransform, as described in the Neural Spline Flow paper (https://arxiv.org/abs/1906.04032).

Code is adapted from the uci.py example from https://github.com/bayesiains/nsf.

A coupling flow fixes half the components of y, and applies a transform to the remaining components, conditioned on the fixed components. This is a restricted form of an autoregressive transform, with a single split into fixed/transformed components.

The transform here is a neural spline flow, where the flow is parametrized by a residual neural network that depends on y_fixed and x. The residual network consists of a sequence of two-layer fully-connected blocks.

:param i: int index of transform in sequence :param param_dim: int dimensionality of y :param context_dim: int = None dimensionality of x :param hidden_dim: int = 512 number of hidden units per layer :param num_transform_blocks: int = 2 number of transform blocks comprising the transform :param activation: str = ‘relu’ activation function :param dropout_probability: float = 0.0 dropout probability for regularization :param batch_norm: bool = False whether to use batch normalization :param num_bins: int = 8 number of bins for the spline :param tail_bound: float = 1. :param apply_unconditional_transform: bool = False whether to apply an unconditional transform to fixed components :param base_transform_type: str = ‘rq-coupling’ type of base transform, one of {rq-coupling, rq-autoregressive}

:return: Transform the NSF transform