scimba_torch.neural_nets.structure_preserving_nets.invertible_nn¶

An invertible neural network made of RealNVP layers.

Classes

`InvertibleLayer`(size, conditional_size, **kwargs)	An abstract class for an invertible layer.
`InvertibleNet`(size, conditional_size[, ...])	An invertible neural network composed of multiple invertible layers.

class InvertibleLayer(size, conditional_size, **kwargs)[source]¶

Bases: ScimbaModule, ABC

An abstract class for an invertible layer.

Parameters:

size (int) – dimension of the input data
conditional_size (int) – dimension of the conditional input data
**kwargs – other arguments for the layer

abstract backward(inputs, with_last_layer=True)[source]¶

Abstract method for the backward pass of the invertible layer.

Parameters:

inputs (Tensor) – the input tensor
with_last_layer (bool) – whether to use the last layer of the network or not (default: True)

abstract log_abs_det_jacobian(y, mu)[source]¶

Computes the log absolute value of the determinant of the Jacobian.

This method must be implemented by all subclasses. It is the primary method for computing the Jacobian determinant as it is numerically stable.

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 of the Jacobian as a tensor of shape (batch_size,).

abs_det_jacobian(y, mu)[source]¶

Computes the absolute value of the determinant of the Jacobian.

Default implementation uses exp(log_abs_det_jacobian). Subclasses can override this method if they have a more efficient direct computation.

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 of the Jacobian as a tensor of shape (batch_size,).

class InvertibleNet(size, conditional_size, layers_list=None, **kwargs)[source]¶

Bases: ScimbaModule

An invertible neural network composed of multiple invertible layers.

Parameters:

size (int) – dimension of the input data
conditional_size (int) – dimension of the conditional input data
layers_list (list[InvertibleLayer]) – list of invertible layers to compose
**kwargs – other arguments for the invertible layers

forward(inputs, with_last_layer=True)[source]¶

Applies the forward pass of the invertible network.

Parameters:

inputs (Tensor) – the input tensor of shape (batch_size, dim + p_dim).
with_last_layer (bool) – whether to use the last layer of the network or not (default: True)

Return type:

Tensor

Returns:

The output tensor of shape (batch_size, dim + p_dim) after applying all layers.

backward(inputs, with_last_layer=True)[source]¶

Applies the backward pass of the invertible network.

Parameters:

inputs (Tensor) – the input tensor of shape (batch_size, dim + p_dim).
with_last_layer (bool) – whether to use the last layer of the network or not.

Return type:

Tensor

Returns:

The output tensor of shape (batch_size, dim + p_dim) after applying all layers in reverse.

log_abs_det_jacobian(inputs)[source]¶

Computes the log absolute value of the determinant of the Jacobian.

This method is numerically more stable and is commonly used in log-probability computations for normalizing flows.

Parameters:: inputs (Tensor) – the input tensor of shape (batch_size, dim + p_dim).
Return type:: Tensor
Returns:: The log absolute determinant of the Jacobian as a tensor of shape (batch_size,).

abs_det_jacobian(inputs)[source]¶

Computes the absolute value of the determinant of the Jacobian.

This method is useful for change of variables in integrals. Uses the layer’s abs_det_jacobian method if overridden, otherwise falls back to exp(log_abs_det_jacobian).

Parameters:: inputs (Tensor) – the input tensor of shape (batch_size, dim + p_dim).
Return type:: Tensor
Returns:: The absolute determinant of the Jacobian as a tensor of shape (batch_size,).