scimba_torch.neural_nets.structure_preserving_nets.affine_ode_layers

Affine and constant flow layers for invertible networks.

Classes

AffineFlowLayer(size, conditional_size, ...)	Affine flow layer for RealNVP-style transformations.
ConstantFlowLayer(size, conditional_size, ...)	Constant flow layer for NICE-style transformations.

class ConstantFlowLayer(size, conditional_size, split_sizes, split_index, other_indices, **kwargs)

Bases: ODESplittedLayer

Constant flow layer for NICE-style transformations.

This layer creates len(other_indices) neural networks that progressively incorporate information from the other split parts.

Parameters:

• size (int) – dimension of the input part to process (split_sizes[split_index])

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

• split_sizes (list[int]) – list of sizes for all split parts

• split_index (int) – index of the part this layer processes

• other_indices (list[int]) – list of indices of other parts to use as conditioning

• **kwargs – other arguments for the neural networks

forward(y, mu, other_parts, with_last_layer=True)

Forward pass: x_a stays unchanged, others are shifted.

For K=3: y_a = x_a, y_b = x_b + t(x_a, mu), y_c = x_c + t(x_a, y_b, mu)

Parameters:

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

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

• other_parts (list[Tensor]) – list of other tensor parts [x_b, x_c, …]

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

Return type:

Tensor

Returns:

Recombined tensor with y and transformed other_parts

backward(y, mu, other_parts, with_last_layer=True)

Backward pass: inverse transformation.

For K=3: x_c = y_c - t(y_a, y_b, mu), x_b = y_b - t(y_a, mu), x_a = y_a

Parameters:

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

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

• other_parts (list[Tensor]) – list of transformed parts [y_b, y_c, …]

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

Return type:

Tensor

Returns:

Recombined tensor with y and inverse-transformed other_parts

log_abs_det_jacobian(y, mu, other_parts)

Log absolute determinant of Jacobian.

For NICE-style constant flow, the Jacobian is triangular with 1s on diagonal, so det = 1 and log|det| = 0.

Parameters:

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

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

• other_parts (list[Tensor]) – list of other tensor parts

Return type:

Tensor

Returns:

Zeros tensor of shape (batch_size,)

abs_det_jacobian(y, mu, other_parts)

Absolute determinant of Jacobian.

For NICE-style constant flow, det = 1.

Parameters:

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

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

• other_parts (list[Tensor]) – list of other tensor parts

Return type:

Tensor

Returns:

Ones tensor of shape (batch_size,)

class AffineFlowLayer(size, conditional_size, split_sizes, split_index, other_indices, **kwargs)

Bases: ODESplittedLayer

Affine flow layer for RealNVP-style transformations.

This layer applies affine transformations: y = exp(t) ⊙ x + s where ⊙ is element-wise multiplication, t is log-scale and s is translation.

Parameters:

• size (int) – dimension of the input part to process (split_sizes[split_index])

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

• split_sizes (list[int]) – list of sizes for all split parts

• split_index (int) – index of the part this layer processes

• other_indices (list[int]) – list of indices of other parts to use as conditioning

• **kwargs – other arguments for the neural networks

forward(y, mu, other_parts, with_last_layer=True)

Forward pass: x_a stays unchanged, others are affinely transformed.

For K=3: - y_a = x_a - y_b = exp(t(x_a, mu)) ⊙ x_b + s(x_a, mu) - y_c = exp(t(x_a, y_b, mu)) ⊙ x_c + s(x_a, y_b, mu)

Parameters:

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

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

• other_parts (list[Tensor]) – list of other tensor parts [x_b, x_c, …]

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

Return type:

Tensor

Returns:

Recombined tensor with y and transformed other_parts

backward(y, mu, other_parts, with_last_layer=True)

Backward pass: inverse affine transformation.

For K=3: - x_c = (y_c - s(y_a, y_b, mu)) * exp(-t(y_a, y_b, mu)) - x_b = (y_b - s(y_a, mu)) * exp(-t(y_a, mu)) - x_a = y_a

Parameters:

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

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

• other_parts (list[Tensor]) – list of transformed parts [y_b, y_c, …]

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

Return type:

Tensor

Returns:

Recombined tensor with y and inverse-transformed other_parts

log_abs_det_jacobian(y, mu, other_parts)

Log absolute determinant of Jacobian.

For affine transformation y = exp(t) ⊙ x + s: log|det(J)| = sum(t_i) for each network

Parameters:

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

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

• other_parts (list[Tensor]) – list of other tensor parts

Return type:

Tensor

Returns:

Log determinant tensor of shape (batch_size,)

abs_det_jacobian(y, mu, other_parts)

Absolute determinant of Jacobian.

For affine transformation: det = exp(sum(t_i)) for each network

Parameters:

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

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

• other_parts (list[Tensor]) – list of other tensor parts

Return type:

Tensor

Returns:

Determinant tensor of shape (batch_size,)