tfep.nn.transformers.affine.affine_transformer

tfep.nn.transformers.affine.affine_transformer(x, shift, log_scale)[source]

Implement an affine transformer for autoregressive normalizing flows.

This provides a functional API to the AffineTransformer layer. It implements the transformation

\(y_i = exp(a_i) * x_i + b_i\)

where \(a_i\) and \(b_i\) are the log scale and shift parameters of the transformation that are usually generated by a conditioner.

The function returns the transformed feature as a Tensor of shape (batch_size, n_features) and the log absolute determinant of its Jacobian as a Tensor of shape (batch_size,).

Parameters:

x (torch.Tensor) – Input tensor x of shape (batch_size, n_features).
shift (torch.Tensor) – The shift coefficients of shape (batch_size, n_features) (i.e. the b coefficients).
log_scale (torch.Tensor) – The logarithm of the scale coefficients of shape (batch_size, n_features) (i.e. the a coefficients).

Returns:

y (torch.Tensor) – Output tensor of shape (batch_size, n_features).
log_det_J (torch.Tensor) – The logarithm of the absolute value of the determinant of the Jacobian of the transformation with shape (batch_size,).