tf-gnns 0.2.0

A hackable TensorFlow GraphNets library

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).

tf_gnns - A Hackable GraphNets library

A library for easy construction of message-passing networks in TensorFlow Keras. It is largely inspired by this DeepMind paper and the corresponding open-source library (original graph_nets library).

The tf_gnns library has no external dependencies except TensorFlow 2.x (there is no support for TF 1.x graph/session-based computation). It implements alternative design constraints from graph_nets, taking advantage of Keras facilities to build complex models easily and without large drops in performance.

tf_gnns is built to support arbitrary node/edge/global attributes and update functions. A set of utility functions for MLP construction with Keras is also provided (i.e., handling input/output sizes for valid networks), replacing Sonnet.

The main motivation for this library was the absence of a relatively short and efficient implementation of GNNs explicitly created to take advantage of Keras functionality. GNN implementations that take advantage of tensorflow_probability functionality are planned for future releases (such as the one used in [2]).

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.

Torch backend note

For Keras 3 + Torch backend with Triton enabled, this repository is currently tested with:

• torch==2.11.0
• triton==3.6.0 (installed as a dependency of torch 2.11.0)

Recommended setup:

pip install "torch==2.11.0"
KERAS_BACKEND=torch pytest -q tests

If you are using a different Torch/Triton combo and hit import-time crashes in triton / torch._dynamo, pinning to the combination above is the first step.

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.

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