sheaf_mpnn#

class BaseSheafConv[source]

Bases: MessagePassing

Shared base for all NSD and NSP sheaf convolution layers.

Factors out the parameterization and utilities that are identical across every variant in both model families:

W1 / W2 – bilateral stalk transforms (left d*d, right f*f).
sigma – activation function (Tanh).
reset_parameters() – Xavier init for W1, W2, and any map_generator.
_apply_stalk_transform(x) – computes W1 @ x @ W2.
_apply_norm(...) – abstract; each concrete subclass delegates to the appropriate apply_*_norm function from sheaf_mpnn.utils.

Concrete subclasses must implement:: get_map_products(x_feat, edge_index) -> (self_map, cross_map) _apply_norm(self_map, cross_map, edge_index, num_nodes) forward(x_feat, x_stalk, edge_index) -> updated stalk message(...)

Parameters:

stalk_dim (int)
in_channels (int)
hidden_dim (int)
context_dim (int | None, default: None)
add_self_loops (bool, default: False)
add_lp (bool, default: False)
add_hp (bool, default: False)
sparse_learner (bool, default: False)

map_generator: Module

weight_generator: Linear

__init__(stalk_dim, in_channels, hidden_dim, context_dim=None, add_self_loops=False, add_lp=False, add_hp=False, sparse_learner=False)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

stalk_dim (int)
in_channels (int)
hidden_dim (int)
context_dim (int | None, default: None)
add_self_loops (bool, default: False)
add_lp (bool, default: False)
add_hp (bool, default: False)
sparse_learner (bool, default: False)

sigma: Module

reset_parameters()[source]: Resets all learnable parameters of the module.

message(z_dst, z_src, self_map, cross_map)[source]

Builds per-edge sheaf Laplacian messages.

Parameters:

z_dst – Destination-node transformed stalks [E, d, f].
z_src – Source-node transformed stalks [E, d, f].
self_map – Normalized \(F_\mathrm{dst}^\top F_\mathrm{dst}\) per edge [E, d, d].
cross_map – Normalized \(F_\mathrm{dst}^\top F_\mathrm{src}\) per edge [E, d, d].

Returns:

Per-edge messages [E, d, f].

Return type:

Tensor

class DiagonalNSDConv[source]

Bases: BaseNSDConv

Diagonal NSD convolution layer.

Parameters:

stalk_dim (int)
in_channels (int)
hidden_dim (int)
alpha (float, default: 1.0)
context_dim (int | None, default: None)
add_self_loops (bool, default: False)
add_eps (bool, default: True)
dropout (float, default: 0.0)
add_lp (bool, default: False)
add_hp (bool, default: False)
sparse_learner (bool, default: False)

__init__(stalk_dim, in_channels, hidden_dim, alpha=1.0, context_dim=None, add_self_loops=False, add_eps=True, dropout=0.0, add_lp=False, add_hp=False, sparse_learner=False)[source]

Initializes the shared NSD convolution parameters.

Parameters:

stalk_dim (int) – Number of LEARNED stalk dimensions (reference d). With add_lp/add_hp the full stalk handled by the layer is final_d = stalk_dim + add_lp + add_hp and node states have shape [final_d, in_channels].
in_channels (int) – Feature dimension inside each stalk channel (f).
hidden_dim (int) – Unused by the reference-style linear map generator; kept for API compatibility.
alpha (float, default: 1.0) – Initial residual diffusion step size.
context_dim (int | None, default: None) – Width of each node context vector x_feat.
add_self_loops (bool, default: False) – Whether to add self-loops for degree normalization. Defaults to False; the normalization is already (deg+1)-augmented, so self-loops would augment twice.
add_eps (bool, default: True) – If True, adds a learnable per-stalk rescaling eps in R^d (Remark 4, Appendix A.1). The skip term x_stalk is replaced by (1 + eps) * x_stalk, where eps is broadcast across the feature channel dimension. Initialized to zero so default behaviour is the standard NSD update.
dropout (float, default: 0.0) – Dropout probability applied to the diffusion input (the skip path stays un-dropped), matching the reference per-layer dropout before the left/right weights.
add_lp (bool, default: False) – Appends one fixed stalk channel whose Laplacian block is the signless graph Laplacian D+A (sharpening), as in the reference --add_lp.
add_hp (bool, default: False) – Appends one fixed stalk channel whose Laplacian block is the standard graph Laplacian D-A (smoothing), as in the reference --add_hp.
sparse_learner (bool, default: False) – Uses the reference LocalConcatSheafLearnerVariant input reduction: the concatenated pair is viewed as [final_d, 2f] and summed over the stalk axis, shrinking the learner input from 2*final_d*f to 2f.

get_map_products(x_feat, edge_index)[source]

Compute diagonal map products w_dst**2 and w_dst*w_src for all edges.

Parameters:

x_feat – Node context features [N, context_dim].
edge_index – Graph connectivity [2, E].

Returns:

(self_map, cross_map) both [E, d] – element-wise products.

Return type:

tuple

message(z_dst, z_src, self_map, cross_map)[source]

Builds diagonal Laplacian messages.

Parameters:

z_dst (Tensor) – Destination-node transformed stalks [E, d, f].
z_src (Tensor) – Source-node transformed stalks [E, d, f].
self_map (Tensor) – Normalised w_dst**2 per edge [E, d].
cross_map (Tensor) – Normalised w_dst*w_src per edge [E, d].

Returns:

Per-edge Laplacian messages [E, d, f].

Return type:

Tensor

class GeneralNSDConv[source]

Bases: BaseNSDConv

Generalized NSD convolution layer.

Parameters:

stalk_dim (int)
in_channels (int)
hidden_dim (int)
alpha (float, default: 1.0)
context_dim (int | None, default: None)
add_self_loops (bool, default: False)
use_attention (bool, default: False)
add_eps (bool, default: True)
dropout (float, default: 0.0)
add_lp (bool, default: False)
add_hp (bool, default: False)
sparse_learner (bool, default: False)

__init__(stalk_dim, in_channels, hidden_dim, alpha=1.0, context_dim=None, add_self_loops=False, use_attention=False, add_eps=True, dropout=0.0, add_lp=False, add_hp=False, sparse_learner=False)[source]

Initializes the shared NSD convolution parameters.

Parameters:

stalk_dim (int) – Number of LEARNED stalk dimensions (reference d). With add_lp/add_hp the full stalk handled by the layer is final_d = stalk_dim + add_lp + add_hp and node states have shape [final_d, in_channels].
in_channels (int) – Feature dimension inside each stalk channel (f).
hidden_dim (int) – Unused by the reference-style linear map generator; kept for API compatibility.
alpha (float, default: 1.0) – Initial residual diffusion step size.
context_dim (int | None, default: None) – Width of each node context vector x_feat.
add_self_loops (bool, default: False) – Whether to add self-loops for degree normalization. Defaults to False; the normalization is already (deg+1)-augmented, so self-loops would augment twice.
add_eps (bool, default: True) – If True, adds a learnable per-stalk rescaling eps in R^d (Remark 4, Appendix A.1). The skip term x_stalk is replaced by (1 + eps) * x_stalk, where eps is broadcast across the feature channel dimension. Initialized to zero so default behaviour is the standard NSD update.
dropout (float, default: 0.0) – Dropout probability applied to the diffusion input (the skip path stays un-dropped), matching the reference per-layer dropout before the left/right weights.
add_lp (bool, default: False) – Appends one fixed stalk channel whose Laplacian block is the signless graph Laplacian D+A (sharpening), as in the reference --add_lp.
add_hp (bool, default: False) – Appends one fixed stalk channel whose Laplacian block is the standard graph Laplacian D-A (smoothing), as in the reference --add_hp.
sparse_learner (bool, default: False) – Uses the reference LocalConcatSheafLearnerVariant input reduction: the concatenated pair is viewed as [final_d, 2f] and summed over the stalk axis, shrinking the learner input from 2*final_d*f to 2f.
use_attention (bool, default: False)

get_map_products(x_feat, edge_index)[source]

Compute general map products F_dst^T F_dst and F_dst^T F_src for all edges.

Parameters:

x_feat – Node context features [N, context_dim].
edge_index – Graph connectivity [2, E].

Returns:

(self_map, cross_map) both [E, d, d].

Return type:

tuple

class LowRankNSDConv[source]

Bases: BaseNSDConv

Low-rank NSD convolution layer.

Parameterizes each restriction map as F = A @ B^T where A, B in R^{dxr}, bounding the effective rank of each map to at most rank. This gives cheaper parameterization than GeneralNSDConv (2*d*r vs d*d params per map) while keeping the full d-dimensional stalk.

The Laplacian products reduce to:

restriction_maps_dst^T restriction_maps_dst: = right_dst (left_dst^T left_dst) right_dst^T [d, d]
restriction_maps_dst^T restriction_maps_src: = right_dst (left_dst^T left_src) right_src^T [d, d]

Parameters:

stalk_dim (int)
in_channels (int)
hidden_dim (int)
alpha (float, default: 1.0)
context_dim (int | None, default: None)
add_self_loops (bool, default: False)
rank (int, default: 1)
add_eps (bool, default: True)
dropout (float, default: 0.0)
add_lp (bool, default: False)
add_hp (bool, default: False)
sparse_learner (bool, default: False)

__init__(stalk_dim, in_channels, hidden_dim, alpha=1.0, context_dim=None, add_self_loops=False, rank=1, add_eps=True, dropout=0.0, add_lp=False, add_hp=False, sparse_learner=False)[source]

Initializes a low-rank NSD convolution layer.

Parameters:

stalk_dim (int) – Stalk dimension.
in_channels (int) – Feature dimension inside each stalk channel.
hidden_dim (int) – Hidden width of the restriction-map generator MLP.
alpha (float, default: 1.0) – Initial learnable diffusion step size. Defaults to 1.0.
context_dim (int | None, default: None) – Width of x_feat. Defaults to d * in_channels.
add_self_loops (bool, default: False) – Whether to add self-loops for degree normalization. Defaults to False; the normalization is already (deg+1)-augmented, so self-loops would augment twice.
rank (int, default: 1) – Rank of each restriction map F = A @ B^T. Must be positive.
add_eps (bool, default: True) – Learnable per-stalk skip rescaling (Remark 4). Defaults to True.
dropout (float, default: 0.0) – Dropout probability on the diffusion input. Defaults to 0.0.
add_hp (bool, default: False) – Appends one fixed stalk channel diffusing with the standard Laplacian D-A (reference --add_hp).
add_lp (bool, default: False) – Appends one fixed stalk channel diffusing with the signless Laplacian D+A (reference --add_lp).
sparse_learner (bool, default: False) – Stalk-summed map-generator input (reference LocalConcatSheafLearnerVariant). Defaults to False.

get_map_products(x_feat, edge_index)[source]

Compute low-rank map products F_dst^T F_dst and F_dst^T F_src for all edges.

Each map F = left @ right^T so F^T F = right (left^T left) right^T.

Parameters:

x_feat – Node context features [N, context_dim].
edge_index – Graph connectivity [2, E].

Returns:

(self_map, cross_map) both [E, d, d].

Return type:

tuple

class OrthogonalNSDConv[source]

Bases: BaseNSDConv

Orthogonal NSD convolution layer.

Parameters:

stalk_dim (int)
in_channels (int)
hidden_dim (int)
alpha (float, default: 1.0)
context_dim (int | None, default: None)
add_self_loops (bool, default: False)
clamp_val (float, default: 10.0)
use_attention (bool, default: False)
orth_strategy (Literal['cayley', 'fasth', 'fasthpp', 'householder'], default: "cayley")
add_eps (bool, default: True)
dropout (float, default: 0.0)
add_lp (bool, default: False)
add_hp (bool, default: False)
sparse_learner (bool, default: False)
use_edge_weights (bool, default: False)

__init__(stalk_dim, in_channels, hidden_dim, alpha=1.0, context_dim=None, add_self_loops=False, clamp_val=10.0, use_attention=False, orth_strategy='cayley', add_eps=True, dropout=0.0, add_lp=False, add_hp=False, sparse_learner=False, use_edge_weights=False)[source]

Initializes an orthogonal NSD convolution layer.

The map generator outputs parameters for either: (1) entries of a skew-symmetric matrix (cayley), (2) Householder vectors (fasth), or (3) attention-based mappings. All parameterisations produce orthogonal d x d restriction maps.

Parameters:

stalk_dim (int) – Stalk dimension and orthogonal restriction-map matrix size.
in_channels (int) – Feature dimension inside each stalk channel.
hidden_dim (int) – Hidden width of the restriction-map generator MLP.
alpha (float, default: 1.0) – Initial learnable diffusion step size. Defaults to 1.0.
context_dim (int | None, default: None) – Width of x_feat. Defaults to d * in_channels when omitted.
add_self_loops (bool, default: False) – If True, self-loops augment the degree used for normalization. Defaults to False.
clamp_val (float, default: 10.0) – Maximum absolute value for clamping Cayley-transform parameters. Defaults to 10.0.
use_attention (bool, default: False) – If True, uses the attention-based Cayley initialization from main. Defaults to False.
orth_strategy (Literal['cayley', 'fasth', 'fasthpp', 'householder'], default: "cayley") – “cayley” or “fasth”. Defaults to “cayley”.
add_eps (bool, default: True) – If True, adds a learnable per-stalk rescaling eps in R^d to the skip term (Remark 4). Defaults to True.
dropout (float, default: 0.0) – Dropout probability on the diffusion input. Defaults to 0.0.
use_edge_weights (bool, default: False) – Learns a symmetric per-edge scalar (reference EdgeWeightLearner) that rescales the orthogonal transport maps and replaces the combinatorial degree with the weighted one. Defaults to False.
add_hp (bool, default: False) – Appends one fixed stalk channel diffusing with the standard Laplacian D-A (reference --add_hp).
add_lp (bool, default: False) – Appends one fixed stalk channel diffusing with the signless Laplacian D+A (reference --add_lp).
sparse_learner (bool, default: False) – Stalk-summed map-generator input (reference LocalConcatSheafLearnerVariant). Defaults to False.

get_map_products(x_feat, edge_index)[source]

Compute orthogonal cross-map W_dst^T W_src and identity self-map for all edges.

Parameters:

x_feat – Node context features [N, context_dim].
edge_index – Graph connectivity [2, E].

Returns:

(identity, cross_map) where identity is [E, d, d] and cross_map is [E, d, d]. self_map is the identity because W^T W = I.

Return type:

tuple

message(z_dst, z_src, self_map, cross_map)[source]

Builds orthogonal Laplacian messages.

Parameters:

z_dst (Tensor) – Destination-node transformed stalks [E, d, f].
z_src (Tensor) – Source-node transformed stalks [E, d, f].
self_map (Tensor) – Scalar degree inverse D^{-1}_dst broadcast to [E, 1, 1].
cross_map (Tensor) – Normalised D^{-1/2}_dst W_dst^T W_src D^{-1/2}_src per edge [E, d, d].

Returns:

Per-edge Laplacian messages [E, d, f].

Return type:

Tensor

class NSDModel[source]

Bases: Module

End-to-end Neural Sheaf Diffusion (NSD) model.

The wrapper lifts raw node features into stalk features, applies a stack of NSD convolution layers, and decodes the flattened stalk representation back to the requested output dimension.

Parameters:

in_channels (int)
out_channels (int)
stalk_dim (int, default: 4)
hidden_dim (int, default: 16)
num_layers (int, default: 2)
variant (NSDVariant, default: NSDVariant.GENERAL)
alpha (float, default: 1.0)
add_self_loops (bool, default: False)
orth_strategy (Literal['cayley', 'fasth', 'fasthpp', 'householder'], default: "cayley")
rank (int, default: 1)
add_eps (bool, default: True)
input_dropout (float, default: 0.0)
dropout (float, default: 0.0)
normalize_output (bool, default: False)
jknet (bool, default: False)
add_lp (bool, default: False)
add_hp (bool, default: False)
sparse_learner (bool, default: False)
second_linear (bool, default: False)
use_edge_weights (bool, default: False)
learn_alpha (bool, default: True)

__init__(in_channels, out_channels, stalk_dim=4, hidden_dim=16, num_layers=2, variant=NSDVariant.GENERAL, alpha=1.0, add_self_loops=False, orth_strategy='cayley', rank=1, add_eps=True, input_dropout=0.0, dropout=0.0, normalize_output=False, jknet=False, add_lp=False, add_hp=False, sparse_learner=False, second_linear=False, use_edge_weights=False, learn_alpha=True)[source]

Initializes an NSD model for node-level prediction.

Parameters:

in_channels (int) – Number of raw input features per node.
out_channels (int) – Number of output channels per node (e.g. num classes).
stalk_dim (int, default: 4) – Stalk dimension. Each node is represented internally as a matrix with shape [stalk_dim, hidden_dim].
hidden_dim (int, default: 16) – Feature dimension inside each stalk channel. The encoded node state has size d * hidden_dim.
num_layers (int, default: 2) – Number of NSD convolution layers. Must be positive.
variant (NSDVariant, default: NSDVariant.GENERAL) – Restriction-map family. DIAGONAL is cheapest, GENERAL is most expressive, ORTHOGONAL uses orthogonal maps (via Cayley or Householder parameterisation). GENERAL_ATTENTION and ORTHOGONAL_ATTENTION use an attention-based map initialisation.
alpha (float, default: 1.0) – Initial diffusion step size per layer. Defaults to 1.0, matching the reference update at init.
learn_alpha (bool, default: True) – If False, alpha stays fixed at its initial value (Bodnar-exact update: no step-size parameter); if True it is learned. Defaults to True.
add_self_loops (bool, default: False) – If True, self-loops are added to the graph before computing degree normalization in each layer. Defaults to False: the +1/+I degree augmentation now lives inside the normalization itself, so adding self-loops would augment twice.
orth_strategy (Literal['cayley', 'fasth', 'fasthpp', 'householder'], default: "cayley") – Orthogonality strategy for the ORTHOGONAL variant: “cayley”, “fasth”, “fasthpp”, or “householder” (reference orgqr; fasthpp computes the same map natively). Defaults to “cayley”.
rank (int, default: 1) – Rank of each restriction map for the LOW_RANK variant. Must be positive. Ignored for other variants. Defaults to 1.
add_eps (bool, default: True) – If True, each layer applies a learnable per-stalk rescaling eps in R^d to the skip term (Remark 4). Defaults to True (reference parity: the original NSD always gates the skip with 1 + tanh(eps), eps initialized at zero).
input_dropout (float, default: 0.0) – Dropout probability applied to raw input features before encoding. Defaults to 0.0.
dropout (float, default: 0.0) – Dropout probability applied to stalk features between layers. Defaults to 0.0.
normalize_output (bool, default: False) – If True, L2-normalise the representation before the decoder (Lv et al., 2021). If jknet is True, each layer’s output is also normalised before concatenation. Defaults to False (reference parity: the original NSD decodes the raw representation).
jknet (bool, default: False) – If True, collect hidden states from every layer and concatenate them before the decoder (Xu et al., 2018). Normalization is controlled by normalize_output. Intended for link prediction. Defaults to False.
add_lp (bool, default: False) – Appends one fixed stalk channel diffusing with the signless Laplacian D+A (reference --add_lp).
add_hp (bool, default: False) – Appends one fixed stalk channel diffusing with the standard Laplacian D-A (reference --add_hp).
sparse_learner (bool, default: False) – Stalk-summed map-generator input (reference LocalConcatSheafLearnerVariant), shrinking the learner from 2*final_d*hidden to 2*hidden inputs.
second_linear (bool, default: False) – Adds a second encoder linear after the ELU+dropout stage (reference --second_linear).
use_edge_weights (bool, default: False) – Learned symmetric per-edge scalars rescaling orthogonal maps (reference edge_weights). Defaults to False.

reset_parameters()[source]

forward(x, edge_index)[source]

Runs the NSD encoder, diffusion layers, and decoder.

Parameters:

x (Tensor) – Raw node features with shape [num_nodes, in_channels].
edge_index (Tensor) – Graph connectivity in COO format with shape [2, num_edges].

Returns:

Node outputs with shape [num_nodes, out_channels].

Return type:

Tensor

class NSDVariant[source]

Bases: Enum

DIAGONAL = 1

GENERAL = 2

ORTHOGONAL = 3

GENERAL_ATTENTION = 4

ORTHOGONAL_ATTENTION = 5

LOW_RANK = 6

property layer_class

property layer_kwargs: dict[str, Any]

build_kwargs(orth_strategy='cayley', rank=1, use_edge_weights=False)[source]

Parameters:

orth_strategy (Literal['cayley', 'fasth', 'fasthpp', 'householder'], default: "cayley")
rank (int, default: 1)
use_edge_weights (bool, default: False)

Return type:

dict[str, Any]

NSD layers#

`DiagonalNSDConv`	Diagonal NSD convolution layer.
`GeneralNSDConv`	Generalized NSD convolution layer.
`OrthogonalNSDConv`	Orthogonal NSD convolution layer.
`LowRankNSDConv`	Low-rank NSD convolution layer.

NSD model#

`NSDModel`	End-to-end Neural Sheaf Diffusion (NSD) model.
`NSDVariant`