GNNs need Sparse Tensors; can we do this?

[santi-mir]

Hi everyone,

Motivation

I describe very briefly what I'd like to do (related to OP.), question is at the bottom.

Here is a very well cited and simple paper (40K) for convolutions on graphs (I'd like to use it for chemistry.)

Torch Geometric

It would be easy with Torch Geometric (PyG):

PyG (PyTorch Geometric) is a library built upon PyTorch to easily write and train Graph Neural Networks (GNNs) for a wide range of applications related to structured data

I want to use Keras for multi backend.

The formula in the paper seems trivial to implement in a layer:

but they talk about linear increase of complexity with the number of edges, surely this is due to sparse matrices.

Question 0: Yet, it seems Sparse Tensors support still on queue?

Question 1: We need to re-implement the graph-like layers to get multi-backend, right ? Otherwise one would just PyG.

Question 2: Can one use such sparse representations in pure Keras, chaging then the backend ? I can't find a tutorial or ref, apart from TF itself here.

[Kxrma47]

You do not necessarily need a sparse-matrix type to keep graph convolution at O(E). A portable representation is an edge list plus node features:

node_features: [num_nodes, channels]
edge_index:    [2, num_edges]   # source, destination

A Keras layer can gather source-node features, compute messages, and aggregate them with keras.ops.segment_sum (or the corresponding scatter operation). That touches each edge once and avoids materializing a dense N x N adjacency matrix.

Conceptually:

src, dst = edge_index[0], edge_index[1]
messages = keras.ops.take(node_features, src, axis=0)
aggregated = keras.ops.segment_sum(
    messages, dst, num_segments=keras.ops.shape(node_features)[0]
)

Then add normalization, self-loops, and a learned projection according to the exact GCN formulation. For batches of graphs, concatenate their nodes and offset each graph's edge indices, while retaining a graph-id vector for pooling.

There is no single sparse tensor abstraction whose storage, autograd, and supported kernels are identical across TensorFlow, JAX, and PyTorch. Using a backend-native sparse tensor therefore gives up some portability. Edge-list message passing built from keras.ops is usually the better multi-backend starting point; backend-specific libraries such as PyG remain preferable when you need their optimized operators, samplers, and large graph ecosystem.