GNNModule

Trait GNNModule 

Source
pub trait GNNModule: Module {
    // Required method
    fn forward_gnn(&self, x: &Tensor, edge_index: &[EdgeIndex]) -> Tensor;
}
Expand description

Trait for GNN modules that process graph-structured data.

Unlike regular Module, GNN layers require graph structure in addition to node features.

Required Methods§

Source

fn forward_gnn(&self, x: &Tensor, edge_index: &[EdgeIndex]) -> Tensor

Forward pass with graph structure.

§Arguments
  • x - Node features [num_nodes, in_features]
  • edge_index - List of edges as (source, target) pairs
§Returns

Updated node features [num_nodes, out_features]

Implementors§