scimba_torch.neural_nets.structure_preserving_nets.coupling_layers

Coupling layer for invertible networks.

Classes

CouplingLayer(size, conditional_size, ...[, ...])

A coupling layer that splits input and applies ODE-based transformations.

class CouplingLayer(size, conditional_size, num_splits, ode_layer_type, seed=None, **kwargs)

Bases: InvertibleLayer

A coupling layer that splits input and applies ODE-based transformations.

This layer:

1. Splits the input y into K parts (K = 2, 3, or 4)

2. Creates a random permutation of indices [0, 1, …, K-1]

3. Applies K ODE_splitted_layer transformations, each processing one part while being conditioned on the other parts

Parameters:

• size (int) – total dimension of the input data

• conditional_size (int) – dimension of the conditional input data

• num_splits (int) – number of splits (K = 2, 3, or 4)

• ode_layer_type (type[ODESplittedLayer]) – class type of ODE_splitted_layer to use

• seed (int) – random seed for the permutation (optional, for reproducibility)

• **kwargs – other arguments for the invertible layer

Raises:

ValueError – If num_splits is not 2, 3, or 4.

Example

>>> layer = CouplingLayer(
...     size=10,
...     conditional_size=5,
...     num_splits=3,
...     ode_layer_type=MyODELayer,
...     ode_layer_kwargs={'hidden_dim': 64}
... )

forward(y, mu, with_last_layer=True)

Forward pass through the coupling layer.

Parameters:

• y (Tensor) – the input tensor of shape (batch_size, size)

• mu (Tensor) – the conditional input tensor of shape (batch_size, conditional_size)

• with_last_layer (bool) – whether to use the last layer

Return type:

Tensor

Returns:

The transformed tensor of shape (batch_size, size)

backward(y, mu, with_last_layer=True)

Backward pass (inverse) through the coupling layer.

Parameters:

• y (Tensor) – the input tensor of shape (batch_size, size)

• mu (Tensor) – the conditional input tensor of shape (batch_size, conditional_size)

• with_last_layer (bool) – whether to use the last layer

Return type:

Tensor

Returns:

The inverse-transformed tensor of shape (batch_size, size)

log_abs_det_jacobian(y, mu)

Computes the log absolute determinant of the Jacobian.

Parameters:

• y (Tensor) – the input tensor of shape (batch_size, size)

• mu (Tensor) – the conditional input tensor of shape (batch_size, conditional_size)

Return type:

Tensor

Returns:

The log absolute determinant as a tensor of shape (batch_size,)

abs_det_jacobian(y, mu)

Computes the absolute determinant of the Jacobian.

This is the product of the determinants from each ODE layer.

Parameters:

• y (Tensor) – the input tensor of shape (batch_size, size)

• mu (Tensor) – the conditional input tensor of shape (batch_size, conditional_size)

Return type:

Tensor

Returns:

The absolute determinant as a tensor of shape (batch_size,)