piecewise_linear module

Implementation of the piecewise linear coupling cell This means that the variable transform is piecewise-linear.

class ElementWisePWLinearTransform(*args, **kwargs)[source]

Bases: zunis.models.flows.coupling_cells.transforms.InvertibleTransform

Invertible piecewise-linear transformations over the unit hypercube

Implements a batched bijective transformation h from the d-dimensional unit hypercube to itself, in an element-wise fashion (each coordinate transformed independently)

In each direction, the bijection is a piecewise-linear transform with b bins where the forward transform has evenly spaced bins. The transformation in each bin is actually an affine transformation. The slopes for each direction and each point in the batch are given by an unormalized tensor q_tilde. This input is softmax-normalized such that 1. h(0) = 0 2. h(1) = 1 3. h is monotonous 4. h is continuous

for which knowing the slopes in each bin is sufficient (when the abuse of language “linear”)

Conditions 1. to 3. ensure the transformation is a bijection and therefore invertible The inverse is also an element-wise, piece-wise linear transformation, but, of course, with variable input bin sizes (and fixed output bin sizes).

Constructor for InvertibleTransform

No arguments, no returns, this is just encapsulation for a function and its inverse

static backward(x, q_tilde, compute_jacobian=True)[source]

Apply an element-wise piecewise-linear transformation to some variables

Parameters

  • x (torch.Tensor) – a tensor with shape (N,k) where N is the batch dimension while k is the dimension of the variable space. This variable span the k-dimensional unit hypercube

  • q_tilde (torch.Tensor) – is a tensor with shape (N,k,b) where b is the number of bins. This contains the un-normalized heights of the bins of the piecewise-constant PDF for dimension k, i.e. q_tilde lives in all of R and we don’t impose a constraint on their sum yet. Normalization is imposed in this function using softmax.

  • compute_jacobian (bool, optional) – determines whether the jacobian should be compute or None is returned

Returns

pair (y,h). - y is a tensor with shape (N,k) living in the k-dimensional unit hypercube - j is the jacobian of the transformation with shape (N,) if compute_jacobian==True, else None.

Return type

tuple of torch.Tensor

static forward(y, q_tilde, compute_jacobian=True)[source]

Apply the inverse of an element-wise piecewise-linear transformation to some variables

Parameters

  • y (torch.Tensor) – a tensor with shape (N,k) where N is the batch dimension while k is the dimension of the variable space. This variable span the k-dimensional unit hypercube

  • q_tilde (torch.Tensor) – is a tensor with shape (N,k,b) where b is the number of bins. This contains the un-normalized heights of the bins of the piecewise-constant PDF for dimension k, i.e. q_tilde lives in all of R and we don’t impose a constraint on their sum yet. Normalization is imposed in this function using softmax.

  • compute_jacobian (bool, optional) – determines whether the jacobian should be compute or None is returned

Returns

pair (x,h). - x is a tensor with shape (N,k) living in the k-dimensional unit hypercube - j is the jacobian of the transformation with shape (N,) if compute_jacobian==True, else None.

Return type

tuple of torch.Tensor

class GeneralPWLinearCoupling(*args, **kwargs)[source]

Bases: zunis.models.flows.coupling_cells.general_coupling.InvertibleCouplingCell

Abstract class implementing a coupling cell based on PW linear transformations

A specific way to predict the parameters of the transform must be implemented in child classes.

Generator for the abstract class GeneralPWLinearCoupling

Parameters

  • d (int) – dimension of the space

  • mask (list of bool) – variable mask which variables are transformed (False) or used as parameters of the transform (True)

training: bool[source]
class PWLinearCoupling(*args, **kwargs)[source]

Bases: zunis.models.flows.coupling_cells.piecewise_coupling.piecewise_linear.GeneralPWLinearCoupling

Piece-wise Linear coupling

Coupling cell using an element-wise piece-wise linear transformation as a change of variables. The transverse neural network is a rectangular dense neural network

Notes

Transformation used: zunis.models.flows.coupling_cells.piecewise_coupling.piecewise_linear.ElementWisePWLinearTransform Neural network used: zunis.models.layers.trainable.ArbitraryShapeRectangularDNN

Generator for PWLinearCoupling

Parameters

  • d (int) –

  • mask (list of bool) – variable mask: which dimension are transformed (False) and which are not (True)

  • n_bins (int) – number of bins in each dimensions

  • d_hidden (int) – dimension of the hidden layers of the DNN

  • n_hidden (int) – number of hidden layers in the DNN

  • input_activation (optional) – pytorch activation function before feeding into the DNN. must be a callable generator without arguments (i.e. a classname or a function)

  • hidden_activation (optional) – pytorch activation function between hidden layers of the DNN. must be a callable generator without arguments (i.e. a classname or a function)

  • output_activation (optional) – pytorch activation function at the output of the DNN. must be a callable generator without arguments (i.e. a classname or a function)

  • use_batch_norm (bool) – whether batch normalization should be used in the DNN.

training: bool[source]
piecewise_linear_inverse_transform(y, q_tilde, compute_jacobian=True)[source]

Apply the inverse of an element-wise piecewise-linear transformation to some variables

Parameters

  • y (torch.Tensor) – a tensor with shape (N,k) where N is the batch dimension while k is the dimension of the variable space. This variable span the k-dimensional unit hypercube

  • q_tilde (torch.Tensor) – is a tensor with shape (N,k,b) where b is the number of bins. This contains the un-normalized heights of the bins of the piecewise-constant PDF for dimension k, i.e. q_tilde lives in all of R and we don’t impose a constraint on their sum yet. Normalization is imposed in this function using softmax.

  • compute_jacobian (bool, optional) – determines whether the jacobian should be compute or None is returned

Returns

pair (x,h). - x is a tensor with shape (N,k) living in the k-dimensional unit hypercube - j is the jacobian of the transformation with shape (N,) if compute_jacobian==True, else None.

Return type

tuple of torch.Tensor

piecewise_linear_transform(x, q_tilde, compute_jacobian=True)[source]

Apply an element-wise piecewise-linear transformation to some variables

Parameters

  • x (torch.Tensor) – a tensor with shape (N,k) where N is the batch dimension while k is the dimension of the variable space. This variable span the k-dimensional unit hypercube

  • q_tilde (torch.Tensor) – is a tensor with shape (N,k,b) where b is the number of bins. This contains the un-normalized heights of the bins of the piecewise-constant PDF for dimension k, i.e. q_tilde lives in all of R and we don’t impose a constraint on their sum yet. Normalization is imposed in this function using softmax.

  • compute_jacobian (bool, optional) – determines whether the jacobian should be compute or None is returned

Returns

pair (y,h). - y is a tensor with shape (N,k) living in the k-dimensional unit hypercube - j is the jacobian of the transformation with shape (N,) if compute_jacobian==True, else None.

Return type

tuple of torch.Tensor