tf-gnns 0.3.0

A hackable TensorFlow GraphNets library

tf_gnns - A Hackable GraphNets library

A library for easy construction of message-passing networks in Keras 3.

It is largely inspired by this DeepMind paper and the corresponding open-source library (original graph_nets library). In addition it contains baseline tested implementations for GCNs.

The tf_gnns library is backend-agnostic through Keras 3. TensorFlow, PyTorch, and JAX backends are tested in CI for core tensor operations and high-level MPNN/GCN smoke tests. The TensorFlow compatibility matrix currently covers TensorFlow 2.17 through 2.21, and the Torch backend smoke matrix covers Torch 2.10, 2.11, and 2.12 CPU wheels.

Initial motivation

This library was initially implemented for GraphNet-style MPNNs and all the other related architectures that can be seen as special cases of Graphnets.

The tf_gnns computations are structured to make them amenable to backend graph compilers such as XLA for TensorFlow/JAX, and torch.compile where applicable.

Note that torch paths are at the moment less performant than torch geometric for eager computation mode.

tf_gnns is built to support arbitrary node/edge/global attributes and update functions.

Philosophy and performance

This tf_gnns framework explicitly avoids containing custom low-level kernels, as framework operations are often more maintainable and well performing.

The tf_gnns computations are formed in a way that allow the graph compilers to create optimized GPU code, that can be reasonably expected to perform as good as the underlying keras-compiled code allows. Often this is close to the limits of the capabilities of the accelerators.

The framework takes advantage of graph computation where available (e.g., jax and tensorflow mainly) for creating fused networks. In preliminary benchmarks speedups with torch.compile were not observed.

A set of utility functions for MLP construction with Keras is also provided (i.e., handling input/output sizes for valid networks), replacing Sonnet.

TensorFlow compatibility and test status

TensorFlow	TensorFlow Probability	Status
2.17.x	0.24.x
2.18.x	0.25.x
2.19.x	0.25.x
2.20.x	0.25.x
2.21.x	0.25.x

The matrix above is validated by scripts/run_tf_matrix_tests.sh and in CI (.github/workflows/tests.yml).

Installing tf_gnns

---

NOTE

The current tested matrix is TensorFlow 2.17 through 2.21 with the matching TensorFlow Probability versions shown above.

---

Install with uv (recommended):

uv sync

Or install with pip:

# optional - recommended:
# pip install tensorflow==2.15 
# pip install tensorflow_probability==0.22
pip install tf_gnns

Run tests:

uv sync --group dev
uv run pytest -v

Run tests with coverage and update badge payload:

scripts/run_coverage.sh

Run compatibility tests across TensorFlow versions:

scripts/run_tf_matrix_tests.sh 2.17 2.18 2.19 2.20 2.21

Execution and compilation

tf_gnns execution paths are eager by default so they can remain backend-portable with Keras 3. If you are using the TensorFlow backend and want graph compilation, compile at the application level:

import tensorflow as tf
from tf_gnns.models.graphnet import GraphNetMLP

model = GraphNetMLP(units=32, core_steps=2)

@tf.function
def train_step(graph_tensor_dict):
    with tf.GradientTape() as tape:
        out = model(graph_tensor_dict)
        loss = tf.reduce_mean(out["nodes"])  # example loss
    grads = tape.gradient(loss, model.trainable_variables)
    return loss, grads

This keeps library internals backend-agnostic while still allowing TensorFlow users to optimize execution.

Backend support note

tf_gnns targets Keras 3 backend portability. The main CI matrix validates:

• TensorFlow backend across TensorFlow 2.17, 2.18, 2.19, 2.20, and 2.21.
• Torch backend smoke tests across Torch 2.10, 2.11, and 2.12 CPU wheels.
• JAX execution through backend-agnostic smoke tests when JAX is available.

For Torch, the CI job installs the selected CPU Torch wheel last and runs with uv run --no-sync so the version under test is not replaced by uv's lockfile sync. GPU-enabled Torch/Triton stacks can be more fragile because Keras, Torch, TensorFlow, and JAX may pull different CUDA dependency versions.

Recommended Torch CPU smoke setup:

pip install --index-url https://download.pytorch.org/whl/cpu "torch==2.12.0"
KERAS_BACKEND=torch pytest -q tests/test_torch_backend_runtime.py tests/test_notebook_torch_backend_flow.py

If you are using a GPU Torch/Triton combo and hit import-time crashes in triton / torch._dynamo, first reproduce with CPU wheels or isolate the backend in a clean environment.

Build the Docker test image for a specific TensorFlow version:

docker build --build-arg TENSORFLOW_VERSION=2.17 -t tf-gnns:test .

Use through Docker

You can build a Docker image that uses tf_gnns with the following command, based on Ubuntu 22:

docker build . -t tf_gnns_215 --network host  --build-arg TENSORFLOW_VERSION=2.15

The container implements some logic to sort out the necessary dependencies. Namely,

• Numpy 1.x is required for tf <= 2.14
• Keras 2 support needs to be enabled for tf >= 2.16
• The tensorflow_probability version is selected through a mapping given the tensorflow version.

Examples

tf_gnns basics

You can inspect some basic functionality in the following Colab notebook:

List sorting example

(Example from the original deepmind/graph_nets library) If you are familiar with the original graph_nets library, this example will help you understand how you can transition to tf_gnns.

Sort a list of elements. This notebook and the accompanying code demonstrates how to use the Graph Nets library to learn to sort a list of elements.

A list of elements is treated as a fully connected graph between the elements. The network is trained to label the start node, and which (directed) edges correspond to the links to the next largest element, for each node.

After training, prediction ability is tested by comparing output to true sorted lists. Then the network's ability to generalize is tested by using it to sort larger lists.

Protein-Protein Interaction example

This example shows how to adapt torch_geometric (aka PyG) inputs to tf_gnns inputs. The notebook can be run end-to-end in Google Colab, and out of the box it gives a test-set F1 score that is competitive with SOTA.

Keras 3 + Torch backend example

This example demonstrates using the higher-level GraphNet constructs with Keras 3 configured for the PyTorch backend.

GCN models

tf_gnns includes sparse GCN implementations for node-classification workloads:

• SparseGCNConv: low-level sparse graph convolution layer.
• SparseGCN: stacked sparse GCN model.
• GCNv2: tunedGNN-style high-level GCN stack with residual paths, normalization, and configurable dropout.

See the OGBN-Arxiv examples:

• notebooks/06_gcn_ogbn_arxiv_tfgnns.ipynb (tf_gnns GCN training workflow, including tunedGNN-style configuration)

Performance

From initial tests, the performance of tf_gnns seems to be at least as good as deepmind/graph_nets when using tensor dictionaries.

Publications using tf_gnns

The library has been used so far in the following publications:

[1] Bayesian graph neural networks for strain-based crack localization

[2] Remaining Useful Life Estimation Under Uncertainty with Causal GraphNets

[3] Relational VAE: A Continuous Latent Variable Model for Graph Structured Data