sheaf_mpnn.nsd.LowRankNSDConv

class LowRankNSDConv

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)

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.

Methods

__init__(stalk_dim, in_channels, hidden_dim)	Initializes a low-rank NSD convolution layer.
add_module(name, module)	Add a child module to the current module.
aggregate(inputs, index[, ptr, dim_size])	Aggregates messages from neighbors as \(\bigoplus_{j \in \mathcal{N}(i)}\).
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
double()	Casts all floating point parameters and buffers to double datatype.
edge_update()	Computes or updates features for each edge in the graph.
edge_updater(edge_index[, size])	The initial call to compute or update features for each edge in the graph.
eval()	Set the module in evaluation mode.
explain_message(inputs, dim_size)
extra_repr()	Return the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(x_feat, x_stalk, edge_index)	Applies one NSD diffusion step to lifted node features.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_map_products(x_feat, edge_index)	Compute low-rank map products F_dst^T F_dst and F_dst^T F_src for all edges.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
ipu([device])	Move all model parameters and buffers to the IPU.
jittable([typing])	Analyzes the MessagePassing instance and produces a new jittable module that can be used in combination with torch.jit.script().
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
message(z_dst, z_src, self_map, cross_map)	Builds per-edge sheaf Laplacian messages.
message_and_aggregate(edge_index)	Fuses computations of message() and aggregate() into a single function.
modules()	Return an iterator over all modules in the network.
mtia([device])	Move all model parameters and buffers to the MTIA.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
parameters([recurse])	Return an iterator over module parameters.
propagate(edge_index[, size])	The initial call to start propagating messages.
register_aggregate_forward_hook(hook)	Registers a forward hook on the module.
register_aggregate_forward_pre_hook(hook)	Registers a forward pre-hook on the module.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_edge_update_forward_hook(hook)	Registers a forward hook on the module.
register_edge_update_forward_pre_hook(hook)	Registers a forward pre-hook on the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post-hook to be run after module's load_state_dict() is called.
register_load_state_dict_pre_hook(hook)	Register a pre-hook to be run before module's load_state_dict() is called.
register_message_and_aggregate_forward_hook(hook)	Registers a forward hook on the module.
register_message_and_aggregate_forward_pre_hook(hook)	Registers a forward pre-hook on the module.
register_message_forward_hook(hook)	Registers a forward hook on the module.
register_message_forward_pre_hook(hook)	Registers a forward pre-hook on the module.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_propagate_forward_hook(hook)	Registers a forward hook on the module.
register_propagate_forward_pre_hook(hook)	Registers a forward pre-hook on the module.
register_state_dict_post_hook(hook)	Register a post-hook for the state_dict() method.
register_state_dict_pre_hook(hook)	Register a pre-hook for the state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
reset_parameters()	Resets all learnable parameters of the module.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
set_submodule(target, module[, strict])	Set the submodule given by target if it exists, otherwise throw an error.
share_memory()	See torch.Tensor.share_memory_().
state_dict(*args[, destination, prefix, ...])
to(*args, **kwargs)
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
update(inputs)	Updates node embeddings in analogy to \(\gamma_{\mathbf{\Theta}}\) for each node \(i \in \mathcal{V}\).
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

Attributes

SUPPORTS_FUSED_EDGE_INDEX
T_destination
call_super_init
decomposed_layers
dump_patches
explain
special_args
map_generator
weight_generator
training

__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)

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)

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