torsh_graph/conv/
gcn.rs

1//! Graph Convolutional Network (GCN) layer implementation
2/// Crate-local result alias: the error type defaults to [`TorshError`],
3/// so both `Result<T>` and `Result<T, OtherError>` stay valid.
4type Result<T, E = torsh_core::error::TorshError> = std::result::Result<T, E>;
5
6use crate::parameter::Parameter;
7use crate::{GraphData, GraphLayer};
8use torsh_tensor::{
9    creation::{randn, zeros},
10    Tensor,
11};
12
13/// Graph Convolutional Network (GCN) layer
14#[derive(Debug)]
15pub struct GCNConv {
16    in_features: usize,
17    out_features: usize,
18    weight: Parameter,
19    bias: Option<Parameter>,
20}
21
22impl GCNConv {
23    /// Create a new GCN convolution layer
24    ///
25    /// # Errors
26    /// Returns an error when the weight or bias tensors cannot be allocated.
27    pub fn new(in_features: usize, out_features: usize, bias: bool) -> Result<Self> {
28        let weight = Parameter::new(randn(&[in_features, out_features])?);
29        let bias = if bias {
30            Some(Parameter::new(zeros(&[out_features])?))
31        } else {
32            None
33        };
34
35        Ok(Self {
36            in_features,
37            out_features,
38            weight,
39            bias,
40        })
41    }
42
43    /// Get input feature dimension
44    pub fn in_features(&self) -> usize {
45        self.in_features
46    }
47
48    /// Get output feature dimension
49    pub fn out_features(&self) -> usize {
50        self.out_features
51    }
52
53    /// Apply graph convolution
54    ///
55    /// Propagates with the Kipf & Welling operator
56    /// `H' = D~^(-1/2) (A + I) D~^(-1/2) X W`, i.e. the renormalized adjacency,
57    /// not the graph Laplacian.
58    ///
59    /// # Errors
60    /// Returns an error when `graph.x` does not have `in_features` columns, or
61    /// when `graph.edge_index` is malformed.
62    pub fn forward(&self, graph: &GraphData) -> Result<GraphData> {
63        let feature_dim = graph.x.shape().dims().get(1).copied().unwrap_or(0);
64        if feature_dim != self.in_features {
65            return Err(torsh_core::error::TorshError::ShapeMismatch {
66                expected: vec![graph.num_nodes, self.in_features],
67                got: graph.x.shape().dims().to_vec(),
68            });
69        }
70
71        // Kipf-Welling propagation operator A_hat = D~^-1/2 (A + I) D~^-1/2
72        let adjacency_hat = crate::utils::gcn_norm(&graph.edge_index, graph.num_nodes)?;
73
74        // Apply graph convolution: A_hat @ X @ W
75        let x_transformed = graph.x.matmul(&self.weight.clone_data())?;
76        let mut output_features = adjacency_hat.matmul(&x_transformed)?;
77
78        // Add bias if present
79        if let Some(ref bias) = self.bias {
80            output_features = output_features.add(&bias.clone_data())?;
81        }
82
83        // Create output graph with transformed features
84        Ok(GraphData {
85            x: output_features,
86            edge_index: graph.edge_index.clone(),
87            edge_attr: graph.edge_attr.clone(),
88            batch: graph.batch.clone(),
89            num_nodes: graph.num_nodes,
90            num_edges: graph.num_edges,
91        })
92    }
93}
94
95impl GraphLayer for GCNConv {
96    fn forward(&self, graph: &GraphData) -> Result<GraphData> {
97        self.forward(graph)
98    }
99
100    fn parameters(&self) -> Vec<Tensor> {
101        let mut params = vec![self.weight.clone_data()];
102        if let Some(ref bias) = self.bias {
103            params.push(bias.clone_data());
104        }
105        params
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use torsh_core::device::DeviceType;
113    use torsh_tensor::creation::from_vec;
114
115    #[test]
116    fn test_gcn_creation() {
117        let gcn = GCNConv::new(8, 16, true).expect("gcn");
118        let params = gcn.parameters();
119        assert_eq!(params.len(), 2); // weight + bias
120    }
121
122    #[test]
123    fn test_gcn_forward() {
124        let gcn = GCNConv::new(3, 8, false).expect("gcn");
125
126        // Create simple test graph
127        let x = from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3], DeviceType::Cpu)
128            .expect("from vec should succeed");
129        let edge_index = from_vec(vec![0.0, 1.0, 1.0, 0.0], &[2, 2], DeviceType::Cpu)
130            .expect("from vec should succeed");
131        let graph = GraphData::new(x, edge_index);
132
133        let output = gcn.forward(&graph).expect("forward");
134        assert_eq!(output.x.shape().dims(), &[2, 8]);
135        assert_eq!(output.num_nodes, 2);
136    }
137}
torsh_graph/conv/gcn.rs

torsh_graph/conv/
gcn.rs
