TensorFlow GNN: a GraphTensor type, a Beam sampler, and a Keras v2 constraint
TensorFlow GNN is a library to build Graph Neural Networks on the TensorFlow platform.
At a glance
- What is it?
- TF-GNN is an Apache-2.0 port of a Google-internal graph learning library. It gives you a heterogeneous graph container, distributed subgraph sampling, and ready-made models, but it pins you to Keras v2 and TensorFlow 2.12 or newer.
- Who is it for?
- Adopt TF-GNN if your data already lives in a TensorFlow pipeline and you need heterogeneous graphs with sampled subgraphs at scale, especially if a Beam job is already part of your stack. Do not adopt it if you are on Keras v3 without a willingness to run the legacy shim, or if you want a framework-neutral graph library.
- Can I use it commercially?
- Yes. Apache-2.0 is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 6 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The gap TF-GNN fills: heterogeneous graphs inside TensorFlow
Most TensorFlow code assumes dense tensors with a fixed shape. A graph does not fit that assumption. Nodes and edges have varying counts, and a realistic graph has several node types and several edge types, each with its own feature set. TF-GNN's answer is tfgnn.GraphTensor, a type that represents a graph with what the README calls a heterogeneous schema, meaning multiple types of nodes and edges. The schema is a first-class concept in the library, documented separately, which tells you the project treats type structure as something you declare rather than something you infer at runtime. The intended audience is not someone fitting a small citation network on a laptop. The README describes the library as an OSS port of a Google-internal library used in a broad variety of contexts, on homogeneous and heterogeneous graphs, and in conjunction with other scalable graph mining tools. That framing matters: the design assumes graphs too large to load whole, which is why a sampler ships with the library rather than being left to the user.
GraphTensor, the Beam sampler, and the runner: how the pieces connect
The library has four visible layers. At the bottom is the GraphTensor type and its schema. Above that sit data preparation tools, most notably a graph sampler that the README describes as converting a huge database into a stream of reasonably-sized subgraphs for training and inference. That is the data flow: a large stored graph goes in, sampled subgraphs come out, and those subgraphs are what the model actually sees. The sampler is documented under a Beam guide, so the distributed path runs through Apache Beam. Above the data layer is a collection of ready-to-use models and Keras layers for building your own. At the top is a high-level API for training orchestration, documented under a runner guide. The separation is deliberate. You can use GraphTensor and the Keras layers without the runner, and you can use the sampler without the prebuilt models. The Encoder/Process/Decoder pattern appears in one of the linked notebooks, which is the standard shape for message-passing models where you encode features, run propagation steps, then decode to a prediction. The README does not claim a specific number of propagation rounds or a default architecture, so treat the models directory as the place to check what is actually provided.
Installation and the Keras v2 requirement that trips people up
The stable release installs with a single command:
pip install tensorflow-gnn
The platform requirements are where the friction lives. TensorFlow 2.12 or higher is required, and the README notes that future releases will raise that floor. Keras v2 is required, as traditionally included with TensorFlow 2.x. The README states plainly that TF-GNN does not work with the new multi-backend Keras v3. Users on TF 2.16 or later must install the compatibility package and set an environment variable:
pip install tf-keras TF_USE_LEGACY_KERAS=1
That is a real constraint, not a footnote. Anyone starting a new project on a current TensorFlow release will hit it on the first import. Apache Beam is needed for distributed graph sampling, so it is optional if you sample locally or feed pre-sampled data. Some tests and scripts that use tensorflow.lite require a separate package, installed with pip install ai-edge-litert. The README also states the library is developed and tested on Linux, and that running on other platforms supported by TensorFlow may be possible. Read that as an untested path. If your team develops on macOS or Windows, you are outside the tested configuration and should expect to spend time on build issues that Linux users never see.
Where TF-GNN is the wrong tool
The Keras v2 pin is the most obvious failure mode, and the README is honest about it. The second is subtler. TF-GNN is a TensorFlow library, so adopting it means adopting TensorFlow's installation surface, its GPU driver requirements, and its release cadence for the life of the project. If your serving stack is not TensorFlow, you are adding a dependency you cannot remove later without rewriting the model. The third limitation is the sampler's dependency on Beam. The README presents the graph sampler as the tool for converting a huge database into subgraphs, and the guide for it is a Beam guide. Beam is a heavyweight distributed runner. If your graph fits in memory, or if you already have a sampling pipeline in another system, the sampler is overhead rather than help. The README does not document alternative sampling backends, so plan around Beam or plan to write your own sampling. Finally, the release history is worth reading. The most recent release listed is v1.0.3 from May 2024, with v1.0.2 in February 2024. The repository's last push is dated September 2026, so development has continued past the last tagged release. Anyone who needs a tagged, versioned artifact should check whether the work since v1.0.3 is in a release yet.
PyTorch Geometric and the difference in approach
PyTorch Geometric is the natural comparison, and the difference is structural rather than cosmetic. PyTorch Geometric builds on PyTorch tensors and its own data loading abstractions, and it is not tied to a single graph container with a declared schema in the way TF-GNN's GraphTensor and schema documentation suggest. TF-GNN's distinguishing move is the combination of a typed heterogeneous container with a Beam-based sampler aimed at graphs too large to materialize. If your work is research-scale, single-machine, and you are already in PyTorch, PyTorch Geometric has the shorter path from idea to running code. If your work is production-scale, heterogeneous, and your training pipeline is TensorFlow, TF-GNN is the one that speaks your stack's language. The honest framing is that the choice follows your framework, not the other way around. Switching frameworks to get a graph library is a large migration, and the README gives no reason to believe TF-GNN is worth that migration on its own.
Licence, maintenance, and what upgrading costs you
The licence is Apache-2.0, which permits commercial use and modification with the usual attribution and notice requirements. This article is not legal advice; read the licence text and your organisation's policy before shipping. On maintenance, the repository is not archived and has recent commits, but the release cadence visible in the supplied material is slow: three tags across roughly a year, ending with v1.0.3 in May 2024. The README warns that future releases will raise the required TensorFlow version, which means upgrades are not free. Each bump in the TensorFlow floor can force a coordinated upgrade of your training environment, your GPU drivers, and your Keras configuration. The Keras v2 to v3 transition is the clearest example: users on TF 2.16 or later need an extra package and an environment variable just to import the library. Budget for that kind of work at every major TensorFlow release rather than assuming the library will absorb it silently. The citation block in the README gives the TF-GNN paper for anyone who needs to reference the design in writing.
Who should start with TF-GNN, and what to check first
Start here if your graphs are heterogeneous, your training runs in TensorFlow, and your data is large enough that sampling subgraphs is part of the design rather than an afterthought. The Beam sampler and the Encoder/Process/Decoder notebooks are the parts worth reading first, because they define the workflow the library expects. Stay away if you are on Keras v3 and cannot run the legacy shim, if your serving stack is PyTorch or JAX, or if your graphs fit comfortably in memory and a full-batch forward pass is fine. Before writing any model code, confirm three things in order: your TensorFlow version is at least 2.12, the tf-keras package plus TF_USE_LEGACY_KERAS=1 resolves your imports on a current release, and the graph sampler's Beam guide covers the source format your data is actually stored in. If that third check fails, you are writing your own sampler, and the value proposition of the library narrows to GraphTensor and the Keras layers alone.
Editorial conclusion
Adopt TF-GNN if your data already lives in a TensorFlow pipeline and you need heterogeneous graphs with sampled subgraphs at scale, especially if a Beam job is already part of your stack. Do not adopt it if you are on Keras v3 without a willingness to run the legacy shim, or if you want a framework-neutral graph library. Before committing, verify three things: that your TensorFlow version is 2.12 or higher, that TF_USE_LEGACY_KERAS=1 and the tf-keras package resolve your Keras import, and that the graph sampler covers your source format, since the README only names Beam as the distributed path.
Community notes