GraphNeuralNetworks.jl Brings Serious GNN Tooling to Julia

Analysis by the aitrendblend editorial team | Graph neural networks | Source paper published in JMLR, April 2025

graph neural networks Julia Flux.jl Lux.jl message passing GPU training

Diagram style illustration of a graph neural network message passing between nodes, representing the GraphNeuralNetworks.jl Julia framework — Message passing between graph nodes, the core operation behind GraphNeuralNetworks.jl and every graph neural network library built after it

Open a Julia REPL, type using GraphNeuralNetworks, and you are three lines away from training a graph convolutional network on a GPU without ever touching Python. That sentence would have sounded like wishful thinking five years ago. Carlo Lucibello and Aurora Rossi’s paper in the Journal of Machine Learning Research lays out how a small Julia team built something that competes, at least on design, with the graph learning tools most researchers already reach for by default.

Key points

GraphNeuralNetworks.jl is not one package but four, GraphNeuralNetworks.jl, GNNLux.jl, GNNlib.jl, and GNNGraphs.jl, split so Flux.jl and Lux.jl users share the same underlying message passing engine.
The framework builds every convolutional layer from two primitives, scatter and gather, wired through apply_edges and reduce_neighborhood, with a fused propagate path for speed.
It handles homogeneous graphs, heterogeneous graphs with multiple node and edge types, and temporal graphs with changing structure over time, all inside one API.
GPU support currently covers CUDA and AMDGPU, with Apple Silicon and TPU backends and multi GPU training named as future work rather than shipped features.
The authors are candid that GraphNeuralNetworks.jl still trails PyTorch Geometric and DGL on large graph training, a limitation tied to the broader Julia deep learning ecosystem rather than this package alone.

The problem Julia keeps trying to solve

Anyone who has moved a model from a Jupyter notebook into a production service knows the pattern. You prototype in a friendly, slow language, then you rewrite the hot path in something faster once it works. Julia was built explicitly to remove that step, a single language fast enough for the inner loop and expressive enough for the outer research code. Graph neural networks are a good stress test for that promise, because the inner loop is unusually irregular. Unlike a convolution over a fixed grid of pixels, a graph convolution has to gather information from a different number of neighbors at every node, then scatter the results back. That kind of variable size computation is exactly where naive implementations fall apart and where a language with real just in time compilation can either shine or embarrass itself.

GraphNeuralNetworks.jl is the answer this particular research group has been building for several years, and this paper is less a new method than a status report and an invitation. There is no new architecture here, no benchmark chasing a leaderboard. It is a software paper, the kind JMLR publishes specifically to give tooling like this the same peer reviewed home as a modeling paper, and that distinction matters for how you should read the rest of this piece. The interesting question is not whether the results beat some baseline. It is whether the design choices hold up, and whether a researcher who already lives in Julia now has a real reason to stay there for graph learning instead of switching to Python the moment graphs show up in a project.

Why this matters Most graph neural network tooling coverage assumes PyTorch Geometric is the only serious option. GraphNeuralNetworks.jl is a working counterexample, actively maintained, GPU capable, and built around a message passing core that generalizes to heterogeneous and temporal graphs without bolting on separate libraries for each.

What the package actually ships

The headline detail, and one easy to miss if you only skim the abstract, is that GraphNeuralNetworks.jl is not really one package. It is four, all living in a single GitHub repository, each with a distinct job.

GNNGraphs.jl, the data layer

Every graph in the ecosystem is represented through a type called GNNGraph, which can store its topology as COO format edge lists, a sparse adjacency matrix, or a dense adjacency matrix depending on what fits the problem. It supports features attached at the node level, the edge level, and the whole graph level, and it can batch many small graphs into one structure for efficient training, the same trick that makes molecule level or social network level datasets tractable. Because it plugs into Graphs.jl, the general purpose Julia graph library, GNNGraph inherits a large existing toolkit rather than reinventing graph algorithms from scratch. Two extended types round this out. GNNHeteroGraph handles graphs where nodes and edges come in multiple types, useful for anything like a citation network with authors, papers, and venues as separate entities. TemporalSnapshotsGNNGraph handles graphs whose structure or features change across a sequence of snapshots, which is the shape most real world traffic, transaction, or contact networks actually take.

GNNlib.jl, the shared engine

This is where the actual message passing math lives, implemented once and reused by both Flux.jl and Lux.jl users so the two front ends never drift apart in behavior. It defines the functional versions of the graph convolutional layers and the core primitives the whole system is built on.

GraphNeuralNetworks.jl and GNNLux.jl, the two front doors

Julia’s deep learning ecosystem currently has two competing philosophies for how a neural network should hold its parameters. Flux.jl favors stateful layers that carry their own weights, the style most people coming from PyTorch will recognize immediately. Lux.jl favors an explicit, functional style where parameters and layer behavior are kept separate, closer in spirit to JAX. Rather than pick a side, the authors built one interface for each. Flux users install GraphNeuralNetworks.jl. Lux users install GNNLux.jl. Both sit on top of the same GNNlib.jl and GNNGraphs.jl foundation, so switching between them is less painful than switching between two unrelated libraries would be.

What this replaced GNNs.jl grew out of GeometricFlux.jl, the first graph neural network library for Julia, which the authors describe as no longer actively maintained. The rewrite brought a cleaner interface, better performance, broader GPU coverage, and a wider set of features, essentially a second generation attempt at the same goal.

How the message passing core actually works

Strip away the Julia specific packaging and the mathematical idea underneath GraphNeuralNetworks.jl is the same message passing neural network framework Gilmer and colleagues formalized back in 2017 for modeling molecules. Every node in a graph updates itself by collecting messages from its neighbors, combining them, and folding that summary into its own new representation.

\[ m_{ij} = \phi\left(x_i, x_j, e_{ij}\right) \] \[ \bar{m}_i = \bigoplus_{j \in \mathcal{N}(i)} m_{ij} \] \[ x_i’ = \psi\left(x_i, \bar{m}_i\right) \]

Here \(x_i\) and \(x_j\) are the features of a node and one of its neighbors, \(e_{ij}\) is whatever information lives on the edge connecting them, and \(\phi\) is a learned message function applied to every edge in the graph. The symbol \(\bigoplus\) stands in for a permutation invariant aggregation, typically a sum or a mean, that lets a node absorb messages from an arbitrary number of neighbors without caring about the order they arrive in. Finally \(\psi\) folds the aggregated message back together with the node’s own previous state to produce its update. Stack a handful of these layers and a node’s representation eventually encodes information from several hops away in the graph.

GraphNeuralNetworks.jl exposes exactly this structure as two primitives, apply_edges, which computes \(m_{ij}\) across every edge, and reduce_neighborhood, which computes the aggregation step. A propagate function fuses the two when it can, turning what would otherwise be a slow per edge loop into something closer to a generalized matrix multiplication. That fusion is where Julia’s multiple dispatch earns its keep, since the compiler can specialize the fused operation differently depending on which message function and aggregation the user picked, without the user writing any of that specialization by hand. The library also leans on KernelAbstractions.jl so a single high level implementation of scatter and gather can be compiled for CPU, CUDA, or AMDGPU targets from one source, instead of maintaining separate kernels per backend.

The paper’s own example of a custom GraphConv layer shows how little boilerplate this leaves for a user who wants to define something new. A struct holds two weight matrices, a bias, an activation function, and an aggregation choice. The forward pass calls propagate with a simple copy of the neighbor features as the message, applies the two weight matrices, and runs the activation. That is a full custom graph convolution in roughly ten lines, and it composes cleanly with the built in layers the package already ships, including graph attention networks and graph isomorphism networks.

“GNNs.jl has quickly established itself as a versatile and comprehensive suite of packages for building and training graph neural networks in Julia.” Lucibello and Rossi, JMLR 2025

Where it sits next to the Python tooling

The authors are refreshingly direct about the comparison instead of dodging it. GraphNeuralNetworks.jl draws ideas from PyTorch Geometric, DGL, and PyTorch Geometric Temporal, and the paper states plainly that it lags behind those Python alternatives in certain areas, large graph training named specifically, and attributes part of that gap to the broader Julia deep learning ecosystem still catching up rather than to any single design flaw in this package.

Aspect	GraphNeuralNetworks.jl	PyTorch Geometric and DGL
Host language	Julia, one language for prototyping and performance	Python with a compiled backend, two language workflow
Front end frameworks	Flux.jl (stateful) and Lux.jl (explicit and functional), sharing one core	PyTorch, generally one dominant style
Graph types supported	Homogeneous, heterogeneous, and temporal in one API	Broad coverage, often via separate extension libraries such as PyTorch Geometric Temporal
GPU backends	CUDA and AMDGPU today, Apple Silicon and TPU planned	CUDA is mature, wider third party backend support overall
Large graph training	Behind, by the authors’ own account	More mature tooling for sampling and scaling to very large graphs
Ecosystem maturity	Younger, smaller community, evolving rapidly	Larger community, more tutorials, more third party layers

None of this makes GraphNeuralNetworks.jl a toy. It makes it a reasonable choice for a specific kind of user, someone already committed to Julia for numerical work, simulation, or scientific computing, who does not want to maintain a second language just for the graph learning piece of a project. For someone starting from zero with no existing Julia investment, the calculus is different, and the honest answer from the paper itself is that Python still has the edge on scale.

What this means if you are choosing a stack

The practical takeaway is less about GraphNeuralNetworks.jl beating anything and more about it closing an option that used to not really exist. A computational biology lab running simulations in Julia, or a physics group already using DifferentialEquations.jl, previously had to either drop into Python for the graph learning component of a pipeline or accept a much rougher tool in GeometricFlux.jl. Now there is a maintained, reasonably complete alternative that plugs into the rest of the Julia scientific stack, including Graphs.jl for graph algorithms and MLDatasets.jl for standard benchmark data, without leaving the language.

The four package split is also a quieter but genuinely useful design decision. Because the message passing math lives in GNNlib.jl separately from either front end, the Flux.jl and Lux.jl interfaces cannot silently diverge in how a given layer behaves, a problem that has bitten other ecosystems where two front ends reimplement the same idea twice and drift apart over time. That kind of structural discipline is easy to overlook in a paper this short, but it is the sort of decision that determines whether a library is still trustworthy in three years or has quietly forked into two incompatible dialects.

A gap worth watching The paper lists incorporating XLA compilation through Reactant.jl, multi GPU training, and Apple Silicon and TPU backends as future plans rather than current features. Anyone evaluating this for a production workload should treat those as roadmap items, not shipped capability, and check the GitHub repository for current status before committing.

Honest limitations

This is a software description, not an empirical study, so there are no accuracy numbers or ablations to interrogate the way there would be for a modeling paper. That is a limitation of the paper itself worth naming plainly. Readers get design rationale and a code example, not benchmark evidence that the fused scatter and gather operations actually deliver a measurable speed advantage over a naive implementation, and not a head to head timing comparison against PyTorch Geometric on a shared task. The claim that GraphNeuralNetworks.jl lags behind on large graph training is stated by the authors themselves, which is a point in favor of the paper’s honesty, but it also means the size of that gap is not quantified anywhere in the text.

The GPU story is narrower than it might first appear too. Only CUDA and AMDGPU are supported at the time of writing, which covers NVIDIA and AMD hardware but leaves Apple Silicon and TPU users waiting on work that has not shipped yet. And because the whole project depends on the pace of Julia’s broader machine learning ecosystem, including how quickly Flux.jl, Lux.jl, and KernelAbstractions.jl themselves mature, GraphNeuralNetworks.jl’s ceiling is partly out of its own maintainers’ hands.

Complete PyTorch reimplementation for illustration

GraphNeuralNetworks.jl is a Julia library, so there is no PyTorch code to lift from the paper itself. What follows is an independent PyTorch reimplementation of the same message passing idea the paper describes, a custom GraphConv layer built from scatter style aggregation, wired into a small model and trained on synthetic graph data, so a reader who works in Python can see the identical concept in a framework they already know.

# Independent PyTorch reimplementation of the message passing pattern
# described in GraphNeuralNetworks.jl (Lucibello and Rossi, JMLR 2025).
# This is not the paper’s own code. It is a from scratch illustration
# of the same scatter and gather idea, built for readers working in PyTorch.

import torch
import torch.nn as nn
import torch.nn.functional as F


class GraphConv(nn.Module):
    “””A minimal message passing graph convolution.

    Mirrors the apply_edges plus reduce_neighborhood pattern from
    GraphNeuralNetworks.jl. Messages are simply the neighbor features,
    aggregated by mean, then combined with the node’s own features
    through two separate learned linear maps.
    “””

    def __init__(self, in_dim, out_dim, aggr=“mean”):
        super().__init__()
        self.lin_self = nn.Linear(in_dim, out_dim)
        self.lin_neigh = nn.Linear(in_dim, out_dim)
        self.aggr = aggr

    def forward(self, x, edge_index):
        # edge_index has shape [2, num_edges], row 0 is source, row 1 is target
        src, dst = edge_index
        messages = x[src]                          # apply_edges equivalent
        out = torch.zeros(x.size(0), messages.size(1), device=x.device)
        out.index_add_(0, dst, messages)            # scatter add
        if self.aggr == “mean”:
            deg = torch.zeros(x.size(0), device=x.device)
            deg.index_add_(0, dst, torch.ones(src.size(0), device=x.device))
            deg = deg.clamp(min=1).unsqueeze(1)
            out = out / deg                         # reduce_neighborhood equivalent
        return F.relu(self.lin_self(x) + self.lin_neigh(out))


class GNNModel(nn.Module):
    “””Two GraphConv layers followed by mean pooling and a linear readout,

    the same shape as the regression example in the paper, translated
    from Flux.jl and GNNChain into plain PyTorch.
    “””

    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        self.conv1 = GraphConv(in_dim, hidden_dim)
        self.conv2 = GraphConv(hidden_dim, hidden_dim)
        self.readout = nn.Linear(hidden_dim, out_dim)

    def forward(self, x, edge_index, batch):
        h = self.conv1(x, edge_index)
        h = self.conv2(h, edge_index)
        num_graphs = int(batch.max().item()) + 1
        pooled = torch.zeros(num_graphs, h.size(1), device=h.device)
        counts = torch.zeros(num_graphs, device=h.device)
        pooled.index_add_(0, batch, h)
        counts.index_add_(0, batch, torch.ones(h.size(0), device=h.device))
        pooled = pooled / counts.clamp(min=1).unsqueeze(1)
        return self.readout(pooled).squeeze(-1)


def make_random_graph(num_nodes=10, num_edges=20, feat_dim=16):
    # Random graph generator, standing in for rand_graph in the paper’s example
    edge_index = torch.randint(0, num_nodes, (2, num_edges))
    x = torch.randn(num_nodes, feat_dim)
    y = torch.randn(1)
    return x, edge_index, y


def collate(graphs):
    # Batches several small graphs into one disjoint union, the same
    # concept as the DataLoader batching shown in the original paper
    xs, edge_indices, ys, batch = [], [], [], []
    node_offset = 0
    for i, (x, edge_index, y) in enumerate(graphs):
        xs.append(x)
        edge_indices.append(edge_index + node_offset)
        ys.append(y)
        batch.append(torch.full((x.size(0),), i, dtype=torch.long))
        node_offset += x.size(0)
    return torch.cat(xs), torch.cat(edge_indices, dim=1), torch.cat(ys), torch.cat(batch)


def evaluate(model, x, edge_index, y, batch):
    model.eval()
    with torch.no_grad():
        pred = model(x, edge_index, batch)
        loss = F.mse_loss(pred, y)
    return loss.item()


def smoke_test():
    torch.manual_seed(0)
    model = GNNModel(in_dim=16, hidden_dim=32, out_dim=1)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    graphs = [make_random_graph() for _ in range(8)]
    x, edge_index, y, batch = collate(graphs)

    model.train()
    for epoch in range(5):
        optimizer.zero_grad()
        pred = model(x, edge_index, batch)
        loss = F.mse_loss(pred, y)
        loss.backward()
        optimizer.step()
        print(f”epoch {epoch} loss {loss.item():.4f}”)

    final_loss = evaluate(model, x, edge_index, y, batch)
    print(f”eval loss {final_loss:.4f}”)
    assert final_loss == final_loss, “loss became NaN, smoke test failed”


if __name__ == “__main__”:
    smoke_test()
  

Run as is, this trains for five steps on random data and prints a decreasing loss, confirming the layer, the batching, and the pooling all fit together without shape errors. It is not a claim that this matches GraphNeuralNetworks.jl’s own performance characteristics, only that the underlying scatter and gather idea the paper describes is portable and easy to verify outside Julia.

Conclusion

The core achievement of this paper is unglamorous in the best sense. Four packages, one shared message passing engine, two front ends for two competing philosophies of how a neural network should hold its state, and enough graph type coverage to handle the messy heterogeneous and time varying data that real applications actually produce instead of the clean single type graphs most tutorials use. That is not a research breakthrough. It is infrastructure, and infrastructure rarely gets the attention it deserves until someone tries to build something on top of it and finds the ground solid or missing.

The conceptual shift worth sitting with is what it means for a scientific computing language to take graph learning seriously as a first class citizen rather than a bolted on extension. Julia’s whole pitch has always been that you should not have to choose between a language that is pleasant to write research code in and one that is fast enough to run it. GraphNeuralNetworks.jl applies that pitch to a domain, graph neural networks, where the irregular, variable degree computation involved has historically punished languages without either careful hand tuned kernels or a mature just in time compiler. Whether that pitch holds up under real workloads is still an open question the paper does not fully answer, but the architecture is at least built to make the attempt honestly.

Transferability beyond Julia’s existing user base is the harder sell. A machine learning engineer with no reason to leave PyTorch is not going to switch languages for graph convolutions alone, and the paper does not pretend otherwise. Where this tool actually earns its place is at the intersection of graph learning and the fields Julia already dominates, differential equations, scientific simulation, computational biology, and physics informed modeling, where staying in one language for an entire pipeline has real value that a marginally more mature Python library does not offset.

The honest remaining limitations are the ones the authors themselves name. Large graph training lags Python alternatives. GPU backend coverage is narrower today than what CUDA focused PyTorch users are used to. And the whole effort rides on the pace of a still maturing Julia deep learning ecosystem that the authors do not fully control. None of that undercuts the value of what has been built so far, it just means GraphNeuralNetworks.jl is a solid foundation rather than a finished cathedral.

What happens next probably matters more than what exists today. Reactant.jl based XLA compilation, multi GPU training, and broader hardware backend support are all named as future work, and each of those would close a real gap rather than add a cosmetic feature. A Julia based graph learning stack that can train on multiple GPUs and compile through XLA would be a genuinely different proposition than the one described in this paper. For now, GraphNeuralNetworks.jl is a capable, well organized tool for the audience it already serves, and a promising but unproven one for everybody else watching from the Python side of the fence.

Frequently asked questions

What is GraphNeuralNetworks.jl

It is an open source suite of four Julia packages for building and training graph neural networks, covering data structures for graphs, the core message passing math, and two separate front end interfaces for the Flux.jl and Lux.jl deep learning frameworks.

Is GraphNeuralNetworks.jl the same as GNNs.jl

Yes. GNNs.jl is the shorthand the authors use in the paper to refer to the whole four package collection, GraphNeuralNetworks.jl, GNNLux.jl, GNNlib.jl, and GNNGraphs.jl, all hosted in one GitHub repository.

Does GraphNeuralNetworks.jl support GPU training

Yes, through the JuliaGPU ecosystem, with CUDA and AMDGPU as the supported backends at the time the paper was published. Apple Silicon and TPU support are listed as future plans rather than current features.

How does it compare to PyTorch Geometric

It offers a comparable design built around message passing, heterogeneous graphs, and temporal graphs, and it borrows some ideas directly from PyTorch Geometric and DGL. The authors state plainly that it currently lags behind those Python libraries in areas such as large graph training, attributing part of the gap to the broader Julia machine learning ecosystem still maturing.

Do I need to know Flux.jl or Lux.jl to use it

You need one or the other, not both. Flux.jl users install GraphNeuralNetworks.jl. Lux.jl users install GNNLux.jl. Both interfaces share the same underlying GNNlib.jl and GNNGraphs.jl packages, so the core behavior is consistent regardless of which front end you pick.

Can I build a custom graph convolutional layer with it

Yes. The package exposes apply_edges and reduce_neighborhood as the two primitives underlying every layer, plus a fused propagate function, so a custom layer can be defined in a small amount of code while still benefiting from GPU compilation through KernelAbstractions.jl.

Want the full technical detail straight from the source.

Read the paper on JMLR Browse the code on GitHub

Lucibello, C. and Rossi, A. GraphNeuralNetworks.jl. Deep learning on graphs with Julia. Journal of Machine Learning Research, 26, pages 1 to 6, 2025. Published April 2025, submitted December 2024.

This analysis is based on the published paper and an independent evaluation of its claims.

GraphNeuralNetworks.jl Brings Serious GNN Tooling to Julia

Key points

The problem Julia keeps trying to solve

What the package actually ships

GNNGraphs.jl, the data layer

GNNlib.jl, the shared engine

GraphNeuralNetworks.jl and GNNLux.jl, the two front doors

How the message passing core actually works

Where it sits next to the Python tooling

What this means if you are choosing a stack

Honest limitations

Complete PyTorch reimplementation for illustration

Conclusion

Frequently asked questions

What is GraphNeuralNetworks.jl

Is GraphNeuralNetworks.jl the same as GNNs.jl

Does GraphNeuralNetworks.jl support GPU training

How does it compare to PyTorch Geometric

Do I need to know Flux.jl or Lux.jl to use it

Can I build a custom graph convolutional layer with it

Related reading

Leave a Comment Cancel Reply