torsh_graph/conv/
gin.rs

1//! Graph Isomorphism Network (GIN) layer implementation
2//! Based on the paper "How Powerful are Graph Neural Networks?"
3// Framework infrastructure - components designed for future use
4#![allow(dead_code)]
5/// Crate-local result alias: the error type defaults to [`TorshError`],
6/// so both `Result<T>` and `Result<T, OtherError>` stay valid.
7type Result<T, E = torsh_core::error::TorshError> = std::result::Result<T, E>;
8
9use crate::parameter::Parameter;
10use crate::{GraphData, GraphLayer};
11use torsh_tensor::{
12    creation::{randn, zeros},
13    Tensor,
14};
15
16/// Graph Isomorphism Network (GIN) layer
17#[derive(Debug)]
18pub struct GINConv {
19    in_features: usize,
20    out_features: usize,
21    eps: f64,
22    train_eps: bool,
23    eps_param: Option<Parameter>,
24    mlp: Vec<Parameter>, // Simple MLP: Linear -> ReLU -> Linear
25    bias: Option<Parameter>,
26}
27
28impl GINConv {
29    /// Create a new GIN convolution layer
30    pub fn new(
31        in_features: usize,
32        out_features: usize,
33        eps: f64,
34        train_eps: bool,
35        bias: bool,
36    ) -> Result<Self> {
37        let eps_param = if train_eps {
38            Some(Parameter::new(torsh_tensor::creation::tensor_scalar(
39                eps as f32,
40            )?))
41        } else {
42            None
43        };
44
45        // Create a simple 2-layer MLP
46        let hidden_dim = (in_features + out_features) / 2;
47        let mlp = vec![
48            Parameter::new(randn(&[in_features, hidden_dim])?),
49            Parameter::new(randn(&[hidden_dim, out_features])?),
50        ];
51
52        let bias = if bias {
53            Some(Parameter::new(zeros(&[out_features])?))
54        } else {
55            None
56        };
57
58        Ok(Self {
59            in_features,
60            out_features,
61            eps,
62            train_eps,
63            eps_param,
64            mlp,
65            bias,
66        })
67    }
68
69    /// Apply GIN convolution
70    pub fn forward(&self, graph: &GraphData) -> Result<GraphData> {
71        let num_nodes = graph.num_nodes;
72        let edge_flat = graph.edge_index.to_vec()?;
73        let num_edges = edge_flat.len() / 2;
74        let edge_data = vec![
75            edge_flat[0..num_edges].to_vec(),
76            edge_flat[num_edges..].to_vec(),
77        ];
78
79        // Build adjacency list for efficient neighbor aggregation
80        let mut adjacency_list: Vec<Vec<usize>> = vec![Vec::new(); num_nodes];
81        for j in 0..edge_data[0].len() {
82            let src = edge_data[0][j] as usize;
83            let dst = edge_data[1][j] as usize;
84            if src < num_nodes && dst < num_nodes {
85                adjacency_list[dst].push(src);
86            }
87        }
88
89        // Aggregate neighbor features (sum aggregation for GIN)
90        let neighbor_features = zeros(&[num_nodes, self.in_features])?;
91
92        for node in 0..num_nodes {
93            let mut aggregated = zeros(&[self.in_features])?;
94
95            // Sum all neighbor features
96            for &neighbor in &adjacency_list[node] {
97                let neighbor_feat = graph
98                    .x
99                    .slice_tensor(0, neighbor, neighbor + 1)?
100                    .squeeze_tensor(0)?;
101                aggregated = aggregated.add(&neighbor_feat)?;
102            }
103
104            let mut node_slice = neighbor_features.slice_tensor(0, node, node + 1)?;
105            let _ = node_slice.copy_(&aggregated.unsqueeze_tensor(0)?);
106        }
107
108        // Get epsilon value
109        let epsilon = if let Some(ref eps_param) = self.eps_param {
110            eps_param.clone_data().to_vec()?[0] as f64
111        } else {
112            self.eps
113        };
114
115        // Combine self and neighbor features: (1 + eps) * h_i + sum(h_j)
116        let self_weighted = graph.x.mul_scalar((1.0 + epsilon) as f32)?;
117        let combined_features = self_weighted.add(&neighbor_features)?;
118
119        // Apply MLP
120        let mut output = combined_features.matmul(&self.mlp[0].clone_data())?;
121
122        // Apply ReLU activation (using max with zero tensor)
123        let zero_tensor = zeros(output.shape().dims())?;
124        output = output.maximum(&zero_tensor)?;
125
126        // Second layer
127        output = output.matmul(&self.mlp[1].clone_data())?;
128
129        // Add bias if present
130        if let Some(ref bias) = self.bias {
131            output = output.add(&bias.clone_data())?;
132        }
133
134        // Create output graph
135        Ok(GraphData {
136            x: output,
137            edge_index: graph.edge_index.clone(),
138            edge_attr: graph.edge_attr.clone(),
139            batch: graph.batch.clone(),
140            num_nodes: graph.num_nodes,
141            num_edges: graph.num_edges,
142        })
143    }
144}
145
146impl GraphLayer for GINConv {
147    fn forward(&self, graph: &GraphData) -> Result<GraphData> {
148        self.forward(graph)
149    }
150
151    fn parameters(&self) -> Vec<Tensor> {
152        let mut params = vec![self.mlp[0].clone_data(), self.mlp[1].clone_data()];
153
154        if let Some(ref eps_param) = self.eps_param {
155            params.push(eps_param.clone_data());
156        }
157
158        if let Some(ref bias) = self.bias {
159            params.push(bias.clone_data());
160        }
161
162        params
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use torsh_core::device::DeviceType;
170    use torsh_tensor::creation::from_vec;
171
172    #[test]
173    fn test_gin_creation() {
174        let gin = GINConv::new(8, 16, 0.5, true, true);
175        let params = gin.expect("operation should succeed").parameters();
176        assert!(params.len() >= 2); // At least MLP weights
177        assert!(params.len() <= 4); // At most MLP + eps + bias
178    }
179
180    #[test]
181    fn test_gin_forward() {
182        let gin = GINConv::new(4, 6, 0.0, false, false);
183
184        // Create test graph
185        let x = from_vec(
186            vec![
187                1.0, 2.0, 3.0, 4.0, // node 0
188                5.0, 6.0, 7.0, 8.0, // node 1
189                9.0, 10.0, 11.0, 12.0, // node 2
190            ],
191            &[3, 4],
192            DeviceType::Cpu,
193        )
194        .expect("operation should succeed");
195        let edge_index = from_vec(vec![0.0, 1.0, 2.0, 1.0, 2.0, 0.0], &[2, 3], DeviceType::Cpu)
196            .expect("from vec should succeed");
197        let graph = GraphData::new(x, edge_index);
198
199        let output = gin
200            .expect("operation should succeed")
201            .forward(&graph)
202            .expect("operation should succeed");
203        assert_eq!(output.x.shape().dims(), &[3, 6]);
204        assert_eq!(output.num_nodes, 3);
205    }
206
207    #[test]
208    fn test_gin_trainable_eps() {
209        let gin_fixed = GINConv::new(4, 8, 1.0, false, false);
210        let gin_trainable = GINConv::new(4, 8, 1.0, true, false);
211
212        let fixed_params = gin_fixed.expect("operation should succeed").parameters();
213        let trainable_params = gin_trainable
214            .expect("operation should succeed")
215            .parameters();
216
217        // Trainable eps version should have one more parameter
218        assert_eq!(trainable_params.len(), fixed_params.len() + 1);
219    }
220}
torsh_graph/conv/gin.rs

torsh_graph/conv/
gin.rs
