topobench.nn.backbones.graph.nsd_utils.sheaf_base module

class SheafDiffusion(edge_index, args)

Bases: Module

Base class for sheaf diffusion models.

This class provides the foundational structure for all sheaf diffusion variants, storing common parameters and configurations.

Parameters:

edge_indextorch.Tensor or None

Edge indices of shape [2, num_edges]. Can be None for inductive models.

argsdict

Configuration dictionary containing: - d (int): Dimension of the stalk space (must be > 0). - hidden_channels (int): Number of hidden channels per stalk dimension. - device (str): Device to run the model on. - layers (int): Number of diffusion layers. - input_dropout (float): Dropout rate for input layer. - dropout (float): Dropout rate for hidden layers. - input_dim (int): Dimension of input features. - output_dim (int): Dimension of output features. - sheaf_act (str): Activation function for sheaf learning. - orth (str): Orthogonalization method.

__init__(edge_index, args)

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