Operations

See More

Structural operations for combining, splitting, and factorizing features in probabilistic circuits.

Cat

Concatenates multiple modules along the feature or channel dimension.

class spflow.modules.ops.Cat(inputs, dim=-1)

Bases: Module

__init__(inputs, dim=-1)

Initialize concatenation operation.

Parameters:

• inputs (list[Module]) – Modules to concatenate.

• dim (int) – Concatenation dimension (0=batch, 1=feature, 2=channel).

log_likelihood(data, cache=None)

Compute log likelihood by concatenating input log-likelihoods.

Parameters:

• data (Tensor) – Input data tensor.

• cache (Cache | None) – Optional cache for storing intermediate results.

Returns:

Concatenated log-likelihood tensor.

Return type:

Tensor

marginalize(marg_rvs, prune=True, cache=None)

Marginalize out specified random variables.

Parameters:

• marg_rvs (list[int]) – List of random variable indices to marginalize.

• prune (bool) – Whether to prune unnecessary modules after marginalization.

• cache (Cache | None) – Optional cache for storing intermediate results.

Returns:

Marginalized module or None if fully marginalized.

Return type:

Optional[Module]

property feature_to_scope: ndarray

Mapping from output features to their respective scopes.

Returns:

2D-array of scopes. Each row corresponds to an output feature,

each column to a repetition.

Return type:

np.ndarray[Scope]

Split

Abstract base class for feature splitting operations.

class spflow.modules.ops.split.Split(inputs, dim=1, num_splits=2)

Bases: Module, ABC

Abstract base class for tensor splitting operations.

Splits input tensors along specified dimensions. Concrete implementations must provide feature_to_scope property.

inputs

Single input module to split.

Type:

nn.ModuleList

dim

Dimension along which to split (0=batch, 1=feature, 2=channel).

Type:

int

num_splits

Number of splits to create.

Type:

int

scope

Variable scope inherited from input.

Type:

Scope

__init__(inputs, dim=1, num_splits=2)

Initialize split operation.

Parameters:

• inputs (Module) – Input module to split.

• dim (int) – Dimension along which to split (0=batch, 1=feature, 2=channel).

• num_splits (int | None) – Number of parts to split into.

get_out_shapes(event_shape)

Get output shapes for each split based on input event shape.

Parameters:

event_shape – Shape of the input event tensor.

Returns:

List of tuples representing output shapes for each split.

marginalize(marg_rvs, prune=True, cache=None)

Marginalize out specified random variables.

Parameters:

• marg_rvs (list[int]) – List of random variable indices to marginalize.

• prune (bool) – Whether to prune the resulting module.

• cache (Optional[Dict[str, Any]]) – Cache dictionary for intermediate results.

Return type:

None | Module

Returns:

Marginalized module or None if fully marginalized.

abstractmethod merge_split_indices(*split_indices)

Merge per-split channel indices back to original feature layout.

This method takes channel indices for each split and combines them into indices matching the original (unsplit) feature layout. Used by parent modules (like EinsumLayer) during sampling.

Parameters:

*split_indices (Tensor) – Channel index tensors for each split, shape (batch, features_per_split).

Return type:

Tensor

Returns:

Merged indices matching the input module’s feature layout, shape (batch, total_features).

abstractmethod merge_split_tensors(*split_tensors)

Merge per-split feature tensors back to original feature layout.

Like merge_split_indices, but keeps all trailing dimensions.

Parameters:

*split_tensors (Tensor) – Tensors for each split with shape (batch, features_per_split, …).

Return type:

Tensor

Returns:

Tensor with shape (batch, total_features, …).

abstract property feature_to_scope: ndarray

Mapping from output features to their respective scopes.

Returns:

2D-array of scopes. Each row corresponds to an output feature,

each column to a repetition.

Return type:

np.ndarray[Scope]

SplitMode

Factory class for creating split configurations.

class spflow.modules.ops.split.SplitMode(split_type, num_splits=2, indices=None)

Bases: object

Configuration for split operations.

Factory class for creating split configurations. Use the class methods to create split configurations that can be passed to modules.

Example

>>> layer = EinsumLayer(inputs=leaf, num_repetitions=3, out_channels=10, split_mode=SplitMode.interleaved(num_splits=3))
>>> layer = LinsumLayer(inputs=leaf, out_channels=10, split_mode=SplitMode.consecutive(num_splits=2))
>>> layer = LinsumLayer(inputs=leaf, out_channels=10, split_mode=SplitMode.by_index(indices=[[0,1], [2,3]]))

classmethod by_index(indices)

Create a split configuration with explicit feature indices.

Splits features according to specified indices. Each inner list contains the feature indices for that split.

Example

>>> SplitMode.by_index([[0, 1, 4], [2, 3, 5, 6, 7]])
SplitMode.by_index(indices=[[0, 1, 4], [2, 3, 5, 6, 7]])

Parameters:

indices (list[list[int]]) – List of lists specifying feature indices for each split. All features must be covered exactly once.

Return type:

SplitMode

Returns:

SplitMode configuration for index-based splitting.

classmethod consecutive(num_splits=2)

Create a consecutive split configuration.

Splits features into consecutive chunks: [0,1,2,3] -> [0,1], [2,3].

Parameters:

num_splits (int) – Number of parts to split into.

Return type:

SplitMode

Returns:

SplitMode configuration for consecutive splitting.

classmethod interleaved(num_splits=2)

Create an interleaved split configuration.

Splits features using modulo: [0,1,2,3] -> [0,2], [1,3].

Parameters:

num_splits (int) – Number of parts to split into.

Return type:

SplitMode

Returns:

SplitMode configuration for interleaved splitting.

__init__(split_type, num_splits=2, indices=None)

Initialize split mode configuration.

Parameters:

• split_type (str) – Type of split (‘consecutive’, ‘interleaved’, or ‘by_index’).

• num_splits (int) – Number of parts to split into.

• indices (list[list[int]] | None) – For ‘by_index’ type, the feature indices for each split.

create(inputs)

Create a Split module with this configuration.

Parameters:

inputs (Module) – Input module to split.

Return type:

Split

Returns:

Split module configured according to this SplitMode.

property indices: list[list[int]] | None

Feature indices for ‘by_index’ split type.

property num_splits: int

Number of splits.

property split_type: str

Type of split (‘consecutive’, ‘interleaved’, or ‘by_index’).

SplitConsecutive

Splits features into consecutive halves (or n parts).

class spflow.modules.ops.SplitConsecutive(inputs, dim=1, num_splits=2)

Bases: Split

Split operation using consecutive feature distribution.

Splits features into consecutive chunks: feature i goes to split i // (num_features / num_splits).

Example

With num_splits=2: [0,1,2,3] -> [0,1], [2,3] With num_splits=3: [0,1,2,3,4,5] -> [0,1], [2,3], [4,5]

__init__(inputs, dim=1, num_splits=2)

Initialize consecutive split operation.

Parameters:

• inputs (Module) – Input module to split.

• dim (int) – Dimension along which to split (0=batch, 1=feature, 2=channel).

• num_splits (int | None) – Number of splits along the given dimension.

log_likelihood(data, cache=None)

Compute log likelihoods for split outputs.

Parameters:

• data (Tensor) – Input data tensor.

• cache (Cache | None) – Optional cache for storing intermediate computations.

Return type:

list[Tensor]

Returns:

List of log likelihood tensors, one for each split output.

merge_split_indices(*split_indices)

Merge split indices back to original layout (consecutive).

SplitConsecutive splits features consecutively: [0,1,2,3] -> [0,1], [2,3]. So we concatenate: left_indices, right_indices.

Return type:

Tensor

merge_split_tensors(*split_tensors)

Merge split feature tensors back to original layout (consecutive).

Return type:

Tensor

property feature_to_scope: ndarray

Mapping from output features to their respective scopes.

Returns:

2D-array of scopes. Each row corresponds to an output feature,

each column to a repetition.

Return type:

np.ndarray[Scope]

SplitInterleaved

Splits features in alternating fashion.

class spflow.modules.ops.SplitInterleaved(inputs, dim=1, num_splits=2)

Bases: Split

Split operation using interleaved feature distribution.

Distributes features using modulo arithmetic: feature i goes to split i % num_splits. Optimized for common cases (2 and 3 splits).

Example

With num_splits=2: [0,1,2,3] -> [0,2], [1,3] With num_splits=3: [0,1,2,3,4,5] -> [0,3], [1,4], [2,5]

split_masks

Boolean masks for each split.

Type:

list[Tensor]

__init__(inputs, dim=1, num_splits=2)

Initialize interleaved split operation.

Parameters:

• inputs (Module) – Input module to split.

• dim (int) – Dimension along which to split.

• num_splits (int | None) – Number of parts to split into.

log_likelihood(data, cache=None)

Compute log likelihoods for each split.

Parameters:

• data (Tensor) – Input data tensor.

• cache (Cache | None) – Optional cache for storing intermediate results.

Return type:

list[Tensor]

Returns:

List of log likelihood tensors, one for each split.

merge_split_indices(*split_indices)

Merge split indices back to original layout (interleaved).

SplitInterleaved splits features by modulo: [0,1,2,3] -> [0,2], [1,3]. So we interleave: [left[0], right[0], left[1], right[1], …].

Return type:

Tensor

merge_split_tensors(*split_tensors)

Merge split feature tensors back to original layout (interleaved).

Return type:

Tensor

property feature_to_scope: ndarray

Get feature-to-scope mapping for each split.

Returns:

Array mapping features to scopes for each split.

Shape: (num_features_per_split, num_splits, num_repetitions)

Return type:

np.ndarray

SplitByIndex

Splits features according to user-specified indices.

class spflow.modules.ops.SplitByIndex(inputs, indices=None, dim=1)

Bases: Split

Split operation using explicit feature indices.

Allows full control over which features go into which split by specifying exact indices for each split.

Example

With indices=[[0, 1, 4], [2, 3, 5, 6, 7]]: features are split into group 1: [0, 1, 4] and group 2: [2, 3, 5, 6, 7]

indices

List of lists specifying feature indices for each split.

inverse_indices

Tensor mapping original positions to split outputs.

__init__(inputs, indices=None, dim=1)

Initialize index-based split operation.

Parameters:

• inputs (Module) – Input module to split.

• indices (Optional[Sequence[Sequence[int]]]) – List of lists specifying feature indices for each split. Each inner list contains the feature indices for that split. All features must be covered exactly once (no overlap, no gaps).

• dim (int) – Dimension along which to split (0=batch, 1=feature, 2=channel).

Raises:

ValueError – If indices are invalid (overlap, gaps, out of bounds).

log_likelihood(data, cache=None)

Compute log likelihoods for each split.

Parameters:

• data (Tensor) – Input data tensor.

• cache (Cache | None) – Optional cache for storing intermediate results.

Return type:

list[Tensor]

Returns:

List of log likelihood tensors, one for each split.

merge_split_indices(*split_indices)

Merge split indices back to original layout.

Takes channel indices for each split and combines them into indices matching the original (unsplit) feature layout.

Parameters:

*split_indices (Tensor) – Channel index tensors for each split.

Return type:

Tensor

Returns:

Merged indices matching the input module’s feature layout.

merge_split_tensors(*split_tensors)

Merge split feature tensors back to original layout.

Return type:

Tensor

property feature_to_scope: ndarray

Get feature-to-scope mapping for each split.

Returns:

Array mapping features to scopes for each split.

Shape: (num_features_per_split, num_splits, num_repetitions)

Return type:

np.ndarray

property indices: list[list[int]]

Get the feature indices for each split.