Implementing a Graph convolutional layer

[wdoppenberg]

Hi all,

I'm trying to implement a candle version of GCNConv, and am struggling with the message passing & aggregation logic. Obviously I need to use some scatter method, so scatter_add is an obvious candidate to use, however I cannot find out how to use it properly in this context. In short, the output of this layer should be [N, F] with N being # nodes and F being the out channels.

An example:

impl GCNConv {
    fn new(in_channels: usize, out_channels: usize, vb: VarBuilder) -> Result<Self> {
        let lin = linear_no_bias(in_channels, out_channels, vb.pp("lin"))?;
        Ok(Self { lin })
    }

    /// # Parameters
    ///
    /// * `xs`: The node feature vectors as matrix ordered by node ID
    /// * `edge_index`: Represents the graph structure in COO format and has shape [2, num_edges]
    fn forward(&self, xs: &Tensor, edge_index: &Tensor) -> candle_core::Result<Tensor> {
        tracing::trace!("Input shape: {:?}", xs.clone().shape());
        let x = self.lin.forward(xs)?;
        let node_weights_shape = x.shape();
        tracing::trace!("Linear shape: {:?}", node_weights_shape);

        // Message & aggregate
        let src_idx = edge_index.i((0, ..))?;
        let dst_idx = edge_index.i((1, ..))?;
        tracing::trace!("Source index shape: {:?}", src_idx.clone().shape());

        let msg = x.index_select(&src_idx, 0)?;
        tracing::trace!("Message shape: {:?}", msg.clone().shape());

        let init = xs.zeros_like()?;
        tracing::trace!("Init shape: {:?}", init.clone().shape());

        let agg = init.scatter_add(&dst_idx, &msg, 0)?;
        tracing::trace!("Aggregated shape: {:?}", agg.clone().shape());
        // let x = (x + agg)?;

        Ok(x)
    }
}

Yields the following:

2024-05-02T06:26:13.137191Z TRACE enzymes: Input shape: [36, 18]
2024-05-02T06:26:13.137609Z TRACE enzymes: Linear shape: [36, 8]
2024-05-02T06:26:13.137675Z TRACE enzymes: Source index shape: [163]
2024-05-02T06:26:13.137703Z TRACE enzymes: Message shape: [163, 8]
2024-05-02T06:26:13.137709Z TRACE enzymes: Init shape: [36, 18]
Error: shape mismatch in scatter-add (self, src), lhs: [36, 18], rhs: [163, 8]

If someone can point me in the right direction I'd be very grateful :)

[LaurentMazare]

Maybe you could give some examples of what you would expect scatter_add to be doing here or of what you want to achieve that would probably make it easier to understand. Fwiw scatter_add should behave in the same way as it does in PyTorch.
Looking at your code but not understanding the context at all, it feels a bit odd for init to be initialized in the same way as the input xs rather than more similar to x as its last dimension will be in_channels whereas I would think you want out_channels instead?

[wdoppenberg]

Thanks for the timely answer!

Indeed the out_channels needs to be the last dimension, thanks for pointing that out. Unfortunately the problem still persists.

I believe it's possible that I'm misunderstanding the functionality of scatter_add, but what I'm trying to achieve is to have for every (directed) edge in the graph a 'message' (an embedding/output of lin) originating from the source node, which is then subsequently aggregated in each destination node by a sum (ideally this op is configurable later to be mean, max, mul). Meaning the output of this layer is a 'convolution' using the graph's structure and the input features.

The full example: https://gist.github.com/wdoppenberg/a22847d84ab55002e002d7828eb03934

[wdoppenberg]

I seem to have found a way, using index_add I get the output as required. Thanks for your input!

[wdoppenberg]

I've updated the gist