Ops¶

Helpers that are only called from other functions in this module are not listed.

rank3_trace(x)[source]¶

Compute the trace of each matrix in a rank-3 tensor.

Parameters:: x (Tensor) – Input tensor of shape \((B, N, N)\).
Returns:: Vector of shape \((B,)\) containing the trace of each matrix.
Return type:: Tensor

rank3_diag(x)[source]¶

Create a batch of diagonal matrices from a rank-2 tensor.

Parameters:: x (Tensor) – Input tensor of shape \((B, N)\).
Returns:: Batched diagonal matrices of shape \((B, N, N)\).
Return type:: Tensor

dense_to_block_diag(adj_pool: Tensor) → Tuple[Tensor, Tensor][source]¶

Convert dense pooled adjacencies to a block-diagonal sparse format.

Parameters:

adj_pool (Tensor) – Dense pooled adjacency of shape \((B, K, K)\) or \((K, K)\).

Returns:

edge_index (~torch.Tensor): Edge indices of shape \((2, E)\) for the block-diagonal adjacency.
edge_weight (~torch.Tensor): Edge weights of shape \((E,)\).

Return type:

tuple

get_mask_from_dense_s(s: Tensor, batch: Tensor | None = None) → Tensor[source]¶

Build a pooled-supernode validity mask from a dense assignment matrix.

The returned mask has shape \([B, K]\), where each entry is True iff supernode \(k\) has at least one assigned input node in graph \(b\).

Supported dense assignment layouts are \(s \in \mathbb{R}^{B \times N \times K}\) and \(s \in \mathbb{R}^{N \times K}\). When s is 2D and batch is provided, rows in s can belong to multiple graphs according to batch. When s is 2D and batch is None, s is treated as a single graph and the output shape is \([1, K]\).

Parameters:

s (Tensor) – Dense assignment tensor of shape \([N, K]\) or \([B, N, K]\).
batch (Tensor, optional) – Node-to-graph assignment of shape \([N]\) for the 2D case.

Returns:

Boolean validity mask of shape \([B, K]\).

Return type:

Tensor

is_multi_graph_batch(batch: Tensor | None) → bool[source]¶

Return True if batch represents more than one graph.

Parameters:: batch (Tensor, optional) – Node-to-graph assignment vector.
Returns:: True when batch is not None, non-empty, and contains at least two distinct graph ids.
Return type:: bool

build_pooled_batch(batch_size: int, num_supernodes: int, device: device, dtype: dtype = torch.int64) → Tensor[source]¶

Build a pooled batch vector for block-structured pooled outputs.

The returned vector has shape \([B \cdot K]\) and follows: \([0, \ldots, 0, 1, \ldots, 1, \ldots, B-1, \ldots, B-1]\), where each graph id is repeated \(K\) times.

This is the standard layout when each input graph contributes exactly \(K\) pooled supernodes stored contiguously.

apply_dense_node_mask(x: Tensor, mask: Tensor) → Tuple[Tensor, Tensor][source]¶

Flatten dense node features and keep only valid rows from mask.

Parameters:

x (Tensor) – Dense node features of shape [B, N, F].
mask (Tensor) – Boolean validity mask of shape [B, N].

Returns:

x_valid (~torch.Tensor): Valid rows from x with shape [N_valid, F].
batch_valid (~torch.Tensor): Graph ids for each valid row with shape [N_valid].

Return type:

tuple

expand_compacted_rows(x_compact: Tensor, valid_mask: Tensor | None, expected_rows: int) → Tensor[source]¶

Expand a compact row-wise tensor back to a padded dense layout.

This helper takes a compact representation containing only valid rows and a boolean validity mask over the full layout, then reconstructs the full tensor by placing compact rows at valid positions and filling invalid rows with zeros.

Parameters:

x_compact (Tensor) – Compact tensor of shape [N_valid, *].
valid_mask (Tensor, optional) – Boolean mask over the full layout. It is flattened internally and must contain expected_rows entries (for example, shape [B, K] for pooled supernodes).
expected_rows (int) – Number of rows in the reconstructed dense tensor.

Returns:

Padded tensor of shape [expected_rows, *].

Return type:

Tensor

Example

Let B=3 and K=4. Suppose valid pooled slots are:

valid_mask = [[1, 1, 0, 0], [1, 0, 1, 1], [1, 0, 0, 0]].

Then expected_rows = B*K = 12, while a compact feature tensor might have only N_valid = 6 rows:

x_compact.shape = [6, F].

Flattening valid_mask yields valid indices [0, 1, 4, 6, 7, 8]. The output has shape [12, F] with compact rows written at those indices and zeros elsewhere.

is_dense_adj(edge_index: Tensor | SparseTensor) → bool[source]¶

Return True if edge_index looks like a dense adjacency matrix.

Accepts a batched dense tensor of shape \((B, N, N)\) or a single dense adjacency matrix of shape \((N, N)\).

postprocess_adj_pool_dense(adj_pool: Tensor, remove_self_loops: bool = False, degree_norm: bool = False, adj_transpose: bool = False, edge_weight_norm: bool = False) → Tensor[source]¶

Postprocess a batched dense pooled adjacency tensor.

Parameters:

adj_pool (Tensor) – Dense pooled adjacency of shape \((B, K, K)\).
remove_self_loops (bool, optional) – If True, zeroes the diagonal. (default: False)
degree_norm (bool, optional) – If True, applies \(\mathbf{D}^{-1/2}\mathbf{A}\mathbf{D}^{-1/2}\) normalization. (default: False)
adj_transpose (bool, optional) – If True, treats the output as transposed when computing row sums for normalization. (default: False)
edge_weight_norm (bool, optional) – If True, normalizes by the maximum absolute edge weight per graph. (default: False)

Returns:

The postprocessed adjacency tensor of shape \((B, K, K)\).

Return type:

Tensor

postprocess_adj_pool_sparse(edge_index: Tensor, edge_weight: Tensor | None, num_nodes: int, remove_self_loops: bool = False, degree_norm: bool = False, edge_weight_norm: bool = False, batch_pooled: Tensor | None = None) → Tuple[Tensor, Tensor | None][source]¶

Postprocess a sparse pooled adjacency in edge_index format.

Parameters:

edge_index (Tensor) – Edge indices of shape \((2, E)\).
edge_weight (Tensor, optional) – Edge weights of shape \((E,)\). (default: None)
num_nodes (int) – Number of pooled nodes.
remove_self_loops (bool, optional) – If True, removes self-loops. (default: False)
degree_norm (bool, optional) – If True, applies symmetric degree normalization to edge weights. (default: False)
edge_weight_norm (bool, optional) – If True, normalizes edge weights by the maximum absolute value per graph. Requires batch_pooled. (default: False)
batch_pooled (Tensor, optional) – Batch vector for pooled nodes of shape \((N,)\), used for per-graph normalization. (default: None)

Returns:

edge_index (~torch.Tensor): Filtered edge indices.
edge_weight (~torch.Tensor or None): Filtered edge weights.

Return type:

tuple

connectivity_to_edge_index(edge_index: Tensor | SparseTensor, edge_weight: Tensor | None = None) → Tuple[Tensor, Tensor | None][source]¶

Convert sparse connectivity to edge index and optional weights.

Accepts edge_index as a \((2, E)\) tensor, a torch COO sparse tensor, or a torch_sparse.SparseTensor (not dense batched adjacency), and returns a canonical \((2, E)\) edge index plus optional edge weights of shape \((E,)\).

Parameters:

edge_index (Adj) – Graph connectivity as a dense \((2, E)\) tensor, torch COO sparse tensor, or torch_sparse.SparseTensor.
edge_weight (Tensor, optional) – Edge weights of shape \((E,)\) when edge_index is a dense tensor. Ignored for sparse types. (default: None)

Returns:

edge_index (~torch.Tensor): Edge indices of shape \((2, E)\).
edge_weight (~torch.Tensor or None): Edge weights of shape \((E,)\) or None.

Return type:

tuple

connectivity_to_torch_coo(edge_index: Tensor | SparseTensor, edge_weight: Tensor | None = None, num_nodes: int | None = None) → Tensor[source]¶

Convert sparse connectivity to a coalesced torch COO sparse tensor.

Parameters:

edge_index (Adj) – Graph connectivity as a dense \((2, E)\) tensor, torch COO sparse tensor, or torch_sparse.SparseTensor.
edge_weight (Tensor, optional) – Edge weights of shape \((E,)\) when edge_index is a dense \((2, E)\) tensor. (default: None)
num_nodes (int, optional) – Number of nodes. Inferred from edge_index if None. (default: None)

Returns:

A coalesced torch sparse COO tensor of shape \((N, N)\).

Return type:

Tensor

Raises:

ValueError – If edge_index is not a Tensor or torch_sparse.SparseTensor.

connectivity_to_sparsetensor(edge_index: Tensor | SparseTensor, edge_weight: Tensor | None = None, num_nodes: int | None = None)[source]¶

Convert sparse connectivity to a torch_sparse.SparseTensor.

Requires the torch_sparse package. Accepts the same input types as connectivity_to_edge_index() (edge index, torch COO, or SparseTensor; not dense batched adjacency).

Parameters:

edge_index (Adj) – Graph connectivity as a dense \((2, E)\) tensor, torch COO sparse tensor, or torch_sparse.SparseTensor.
edge_weight (Tensor, optional) – Edge weights of shape \((E,)\) when edge_index is a dense tensor. (default: None)
num_nodes (int, optional) – Number of nodes. Inferred if None. (default: None)

Returns:

A torch_sparse.SparseTensor of shape \((N, N)\).

Return type:

SparseTensor

Raises:

ImportError – If torch_sparse is not installed.

negative_edge_sampling(edge_index: Tensor, num_nodes: int | Tuple[int, int] | None = None, num_neg_samples: int | None = None, method: str = 'auto', force_undirected: bool = False) → Tensor[source]¶

Sample random negative edges for a graph from edge_index.

This function supports both standard and bipartite graphs. For bipartite inputs (num_nodes=(num_src, num_dst)), force_undirected is ignored.

Parameters:

edge_index (Tensor) – The edge indices of shape \((2, E)\).
num_nodes (int or Tuple[int, int], optional) – The number of nodes, i.e. max_val + 1 of edge_index. If given as a tuple, then edge_index is interpreted as a bipartite graph with shape (num_src_nodes, num_dst_nodes). (default: None)
num_neg_samples (int, optional) – The (approximate) number of negative samples to return. If set to None, will try to return a negative edge for every positive edge. (default: None)
method (str, optional) – The method to use for negative sampling, i.e. "sparse", "dense", or "auto". This is a memory/runtime trade-off. "sparse" will work on any graph of any size, but it could retrieve a different number of negative samples. "dense" will work only on small graphs since it enumerates all possible edges. "auto" will automatically choose the best method. (default: "auto")
force_undirected (bool, optional) – If set to True, sampled negative edges will be undirected. (default: False)

Returns:

Negative edge indices of shape \((2, E_{neg})\).

Return type:

Tensor

Example

Standard graph:

edge_index = torch.as_tensor([[0, 0, 1, 2], [0, 1, 2, 3]])
neg_edge_index = negative_edge_sampling(edge_index)

Bipartite graph:

edge_index = torch.as_tensor([[0, 0, 1, 1], [0, 1, 2, 3]])
neg_edge_index = negative_edge_sampling(edge_index, num_nodes=(2, 4))

batched_negative_edge_sampling(edge_index: Tensor, batch: Tensor | Tuple[Tensor, Tensor], num_neg_samples: int | None = None, method: str = 'auto', force_undirected: bool = False) → Tensor[source]¶

Sample random negative edges independently for each graph in a batch.

This applies negative_edge_sampling() graph-by-graph using batch, then concatenates all sampled negative edges in the original global node indexing space.

Parameters:

edge_index (Tensor) – The edge indices of shape \((2, E)\).
batch (Tensor or Tuple[Tensor, Tensor]) – Batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node to a specific example. If given as a tuple, then edge_index is interpreted as a bipartite graph connecting two different node types.
num_neg_samples (int, optional) – The number of negative samples to return for each graph in the batch. If set to None, will try to return a negative edge for every positive edge. (default: None)
method (str, optional) – The method to use for negative sampling, i.e. "sparse", "dense", or "auto". (default: "auto")
force_undirected (bool, optional) – If set to True, sampled negative edges will be undirected. (default: False)

Returns:

Concatenated negative edge indices of shape \((2, E_{neg})\) with edges from all graphs.

Return type:

Tensor

Example

Homogeneous mini-batch:

edge_index = torch.as_tensor([[0, 0, 1, 2], [0, 1, 2, 3]])
edge_index = torch.cat([edge_index, edge_index + 4], dim=1)
batch = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1])

neg_edge_index = batched_negative_edge_sampling(edge_index, batch)

Bipartite mini-batch:

edge_index1 = torch.as_tensor([[0, 0, 1, 1], [0, 1, 2, 3]])
edge_index2 = edge_index1 + torch.tensor([[2], [4]])
edge_index3 = edge_index2 + torch.tensor([[2], [4]])
edge_index = torch.cat([edge_index1, edge_index2, edge_index3], dim=1)
src_batch = torch.tensor([0, 0, 1, 1, 2, 2])
dst_batch = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2])

neg_edge_index = batched_negative_edge_sampling(
    edge_index, (src_batch, dst_batch)
)

pseudo_inverse(edge_index: Tensor) → Tuple[Tensor | SparseTensor, Tensor | None][source]¶

Compute the Moore–Penrose pseudo-inverse of the adjacency matrix.

Input can be a dense \((N, N)\) tensor or a torch COO sparse tensor (converted to dense for the computation). Output format matches input: dense input returns dense pseudo-inverse; sparse input returns a coalesced torch COO sparse tensor with small entries zeroed.

Parameters:: edge_index (Tensor) – Adjacency matrix as a dense \((N, N)\) tensor or torch COO sparse tensor.
Returns:: Pseudo-inverse of the adjacency, in the same format as the input (dense or torch COO sparse).
Return type:: Tensor

weighted_degree(index: Tensor, weights: Tensor | None = None, num_nodes: int | None = None) → Tensor[source]¶

Computes the weighted degree of a given one-dimensional index tensor.

Parameters:

index (Tensor) – Index tensor of shape \((E,)\).
weights (Tensor, optional) – Edge weights tensor of shape \((E,)\). (default: None)
num_nodes (int, optional) – The number of nodes, i.e. max_val + 1 of index. (default: None)

Returns:

Degree vector of shape \((N,)\).

Return type:

Tensor

Adds remaining self loops to the adjacency matrix.

This method extends the method add_remaining_self_loops by allowing to pass a SparseTensor or torch COO sparse tensor as input.

Parameters:

edge_index (Tensor or SparseTensor) – The edge indices.
edge_weight (Tensor, optional) – One-dimensional edge weights. (default: None)
fill_value (float, optional) – The fill value of the diagonal. (default: 1.)
num_nodes (int, optional) – The number of nodes, i.e. max_val + 1 of edge_index. (default: None)

check_and_filter_edge_weights(edge_weight: Tensor) → Tensor | None[source]¶

Check and filter edge weights to ensure they are in the correct shape: \([E]\) or \([E, 1]\).

Parameters:: edge_weight (Tensor) – The edge weights tensor.

Compute the \(\delta\)-GCN propagation matrix for heterophilic message passing.

Constructs the \(\delta\)-GCN propagation matrix from MaxCutPool: differentiable feature-aware Maxcut for pooling in graph neural networks (Abate & Bianchi, ICLR 2025).

The propagation matrix is computed as: \(\mathbf{P} = \mathbf{I} - \delta \cdot \mathbf{L}_{sym}\) where \(\mathbf{L}_{sym}\) is the symmetric normalized Laplacian.

As described in the paper, when \(\delta > 1\), this operator favors the realization of non-smooth (high-frequency) signals on the graph, making it particularly suitable for heterophilic graphs and MaxCut optimization where adjacent nodes should have different values.

Parameters:

edge_index (Tensor or SparseTensor) – Graph connectivity in COO format of shape \((2, E)\), as torch COO sparse tensor, or as SparseTensor.
edge_weight (Tensor, optional) – Edge weights of shape \((E,)\). (default: None)
delta (float, optional) – Delta parameter for heterophilic message passing. When \(\delta > 1\), promotes high-frequency (non-smooth) signals. (default: 2.0)
num_nodes (int, optional) – Number of nodes. If None, inferred from edge_index. (default: None)

Returns:

edge_index (Tensor or SparseTensor): Updated connectivity in the same format as input.
edge_weight (Tensor or None): Updated edge weights of shape \((E',)\) if input was Tensor, or None if input was SparseTensor or torch COO sparse tensor.

Return type:

tuple

create_one_hot_tensor(num_nodes, kept_node_tensor, device, dtype=None)[source]¶

Create a one-hot assignment matrix for kept nodes.

Parameters:

num_nodes (int) – Total number of nodes \(N\).
kept_node_tensor (Tensor) – Indices of kept nodes of shape \((K,)\).
device (device) – Device to create the tensor on.
dtype (dtype, optional) – Desired dtype. (default: torch.float32)

Returns:

One-hot matrix of shape \((N, K + 1)\), where column \(0\) denotes “unassigned” and columns \(1..K\) correspond to kept nodes.

Return type:

Tensor

get_assignments(kept_node_indices, edge_index=None, max_iter=5, batch=None, num_nodes=None)[source]¶

Assigns all nodes in a graph to the closest kept nodes (supernodes) using a hierarchical assignment strategy with message passing (torch COO version).

This function implements a graph-aware node assignment algorithm that combines iterative message passing with random assignment fallback. It’s designed to create cluster assignments where each node in the graph is mapped to one of the provided kept nodes (supernodes).

Algorithm Overview:

Initial Assignment: All kept nodes are assigned to themselves.
Iterative Propagation: For max_iter iterations, unassigned nodes are assigned to supernodes by finding the closest supernode through graph message passing.
Random Fallback: Any remaining unassigned nodes are randomly assigned to supernodes (respecting batch boundaries if provided).

This approach ensures that all nodes receive assignments while prioritizing graph connectivity and topology-aware clustering.

Parameters:

kept_node_indices (Tensor or list) – Indices of nodes to keep as supernodes. These nodes will serve as cluster centers. Can be a tensor or list of integers.
edge_index (Tensor, optional) – Graph connectivity in COO format of shape \((2, E)\) where \(E\) is the number of edges. Required when max_iter > 0 for graph-aware assignment. (default: None)
max_iter (int, optional) – Maximum number of message passing iterations. If 0, uses only random assignment. Higher values allow more distant nodes to be assigned through graph connectivity. (default: 5)
batch (Tensor, optional) – Batch assignment vector of shape \((N,)\) indicating which graph each node belongs to. When provided, ensures nodes are only assigned to supernodes within the same graph. (default: None)
num_nodes (int, optional) – Total number of nodes in the graph(s). If None, inferred from edge_index or batch. Must be provided if both edge_index and batch are None. (default: None)

Returns:

Assignment mapping tensor of shape \((2, N)\) where the first row contains the original node indices \([0, 1, ..., N-1]\) and the second row contains the corresponding cluster (supernode) indices. The cluster indices are renumbered to be consecutive starting from \(0\).

Return type:

Tensor

Raises:

ValueError – If num_nodes, batch, and edge_index are all None (cannot determine graph size).
ValueError – If max_iter > 0 but edge_index is None (cannot perform graph-aware assignment).