torsh-graph 0.1.3

Graph neural network components for ToRSh - powered by SciRS2
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

//! Graph Convolutional Network (GCN) layer implementation

use crate::parameter::Parameter;
use crate::{GraphData, GraphLayer};
use torsh_tensor::{
    creation::{randn, zeros},
    Tensor,
};

/// Graph Convolutional Network (GCN) layer
#[derive(Debug)]
pub struct GCNConv {
    in_features: usize,
    out_features: usize,
    weight: Parameter,
    bias: Option<Parameter>,
}

impl GCNConv {
    /// Create a new GCN convolution layer
    pub fn new(in_features: usize, out_features: usize, bias: bool) -> Self {
        let weight = Parameter::new(
            randn(&[in_features, out_features]).expect("failed to create weight tensor"),
        );
        let bias = if bias {
            Some(Parameter::new(
                zeros(&[out_features]).expect("failed to create bias tensor"),
            ))
        } else {
            None
        };

        Self {
            in_features,
            out_features,
            weight,
            bias,
        }
    }

    /// Get input feature dimension
    pub fn in_features(&self) -> usize {
        self.in_features
    }

    /// Get output feature dimension
    pub fn out_features(&self) -> usize {
        self.out_features
    }

    /// Apply graph convolution
    pub fn forward(&self, graph: &GraphData) -> GraphData {
        // Compute normalized Laplacian matrix
        let laplacian = crate::utils::graph_laplacian(&graph.edge_index, graph.num_nodes, true);

        // Apply graph convolution: L @ X @ W
        let x_transformed = graph
            .x
            .matmul(&self.weight.clone_data())
            .expect("operation should succeed");
        let mut output_features = laplacian
            .matmul(&x_transformed)
            .expect("operation should succeed");

        // Add bias if present
        if let Some(ref bias) = self.bias {
            output_features = output_features
                .add(&bias.clone_data())
                .expect("operation should succeed");
        }

        // Create output graph with transformed features
        GraphData {
            x: output_features,
            edge_index: graph.edge_index.clone(),
            edge_attr: graph.edge_attr.clone(),
            batch: graph.batch.clone(),
            num_nodes: graph.num_nodes,
            num_edges: graph.num_edges,
        }
    }
}

impl GraphLayer for GCNConv {
    fn forward(&self, graph: &GraphData) -> GraphData {
        self.forward(graph)
    }

    fn parameters(&self) -> Vec<Tensor> {
        let mut params = vec![self.weight.clone_data()];
        if let Some(ref bias) = self.bias {
            params.push(bias.clone_data());
        }
        params
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use torsh_core::device::DeviceType;
    use torsh_tensor::creation::from_vec;

    #[test]
    fn test_gcn_creation() {
        let gcn = GCNConv::new(8, 16, true);
        let params = gcn.parameters();
        assert_eq!(params.len(), 2); // weight + bias
    }

    #[test]
    fn test_gcn_forward() {
        let gcn = GCNConv::new(3, 8, false);

        // Create simple test graph
        let x = from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3], DeviceType::Cpu)
            .expect("from vec should succeed");
        let edge_index = from_vec(vec![0.0, 1.0, 1.0, 0.0], &[2, 2], DeviceType::Cpu)
            .expect("from vec should succeed");
        let graph = GraphData::new(x, edge_index);

        let output = gcn.forward(&graph);
        assert_eq!(output.x.shape().dims(), &[2, 8]);
        assert_eq!(output.num_nodes, 2);
    }
}