JittorGeometric 2.0 puts graph neural networks on a JIT framework

A graph library that inherits the Jittor execution model

JittorGeometric 2.0 is the graph machine learning layer for the Jittor framework, which comes out of Tsinghua and compiles models at runtime. The library description in the repository says exactly that: a graph machine learning library built on Jittor, with the primary language listed as Python and a Python 3.10 plus badge.

What the framework choice actually changes is the shape of your model code. The quick tour in the README defines a model with a method named execute rather than forward, and it imports the neural network module from jittor directly, not from torch.

from jittor import nn
from jittor_geometric.nn import GCNConv

class GCN(nn.Module):
    def __init__(self, dataset, dropout=0.8):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(in_channels=dataset.num_features, out_channels=256)
        self.conv2 = GCNConv(in_channels=256, out_channels=dataset.num_classes)
        self.dropout = dropout

So the library is not a thin compatibility layer over an existing graph package. It is written against a different framework's conventions, and anyone coming from PyTorch will have to internalize the execute convention, the jt.no_grad() context manager, and the fact that a step is driven by passing a loss to the optimizer rather than by calling backward.

The optimizer call in the quick tour makes that concrete: optimizer.step(loss) is the whole training step, and the loss is computed as a function of the model output.

Coo to CSR and CSC conversion is a first class step

The quick tour preprocesses the graph into two sparse formats before the model ever runs. It calls gcn_norm on the edge index with add_self_loops=True, then converts the normalized graph into CSC and CSR structures under jt.no_grad().

from jittor_geometric.ops import cootocsr,cootocsc
from jittor_geometric.nn.conv.gcn_conv import gcn_norm
edge_index, edge_weight = data.edge_index, data.edge_attr
edge_index, edge_weight = gcn_norm(
                        edge_index, edge_weight,v_num,
                        improved=False, add_self_loops=True)
with jt.no_grad():
   data.csc = cootocsc(edge_index, edge_weight, v_num)
   data.csr = cootocsr(edge_index, edge_weight, v_num)

The forward pass then takes the node features together with both sparse matrices, and the model never touches the edge list again:

def execute(self):
    x, csc, csr = data.x, data.csc, data.csr
    x = nn.relu(self.conv1(x, csc, csr))
    x = nn.dropout(x, self.dropout, is_train=self.training)
    x = self.conv2(x, csc, csr)
    return nn.log_softmax(x, dim=1)

This is the standard message passing contract: normalize the adjacency, materialize it in the format the convolution kernel wants, then pass features plus the graph. The point worth noting for an evaluator is that the C++ layer is not optional decoration. The setup.py package_data lists ops/cpp source files and headers as package data, which means the sparse operators are shipped as compiled extension sources, and the feature list ties the sparse path to CuSparse acceleration. A pure Python fallback is not described anywhere in the README, so GPU kernels are the intended path rather than an optional one.

Forty plus models across five families, each with an example file

The model catalogue is organised by family, and every entry links to a runnable example script under examples/. The classic family covers ChebNet, GCN, GraphSAGE, GAT, SGC, APPNP and GCNII. The spectral family adds GPRGNN, BernNet, ChebNetII, EvenNet and OptBasis. The dynamic family lists JODIE, DyRep, TGN, GraphMixer and Dygformer for event based and temporal graphs.

The molecular family is the most varied, with SchNet, DimeNet, EGNN, Graphormer, SphereNet, Uni-Mol and Transformer-M, and a separate section covers graph self-supervised learning, starting with DGI. The example directory on disk matches that listing and adds a few entries the tables do not show, such as craft_example.py, grace_example.py, link_pred_example.py, mvgrl_example.py, minibatch_example.py, dyrep_v2_example.py, gcn_acl_example.py, and a dist_gcn directory for the distributed example.

That directory listing is worth a second look, because it tells you which of the 2.0 headline features actually have example code attached. The distributed training claim in the release notes has a dedicated dist_gcn folder, and mini-batch support has minibatch_example.py. New in 2.0 are also listed as multi-GPU and multi-node training over MPI, event based dynamic graph processing with parallel execution, an Ascend-GNN for NPU, and new model categories for graph transformers, self-supervised learning and recommendation systems. The repository also has a frontend_demo directory, which suggests the documentation site and the library are developed side by side.

Datasets and a dependency list pinned to exact versions

Dataset loading is part of the library, with a datasets module exposing Planetoid and the README naming OGB and Reddit among the built-in loaders. The quick tour shows the loader taking a root path, a dataset name and a transform, then indexing the dataset to get a single data object.

import os.path as osp
from jittor_geometric.datasets import Planetoid
import jittor_geometric.transforms as T

dataset = 'cora'
path = osp.join(osp.dirname(osp.realpath(__file__)), '.', 'data', dataset)
dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())
data = dataset[0]

The pinning in requirements.txt is the part that deserves attention before you plan a deployment. The list pins ase to 3.24.0, cupy to 13.3.0, jittor to 1.3.9.14, numpy to 1.24.0, pandas to 2.2.3, scipy to 1.15.1, scikit_learn to 1.6.1, sympy to 1.13.3, tqdm to 4.66.4 and setuptools to 69.5.1, alongside tooling entries such as Flask 3.1.0, pywebio 1.8.3, sphinx_rtd_theme 3.0.2, PyMetis 2023.1.1, schnetpack 2.0.0 and huggingface_hub 0.27.1. The setup.py itself declares the distribution name as jittor_geometric at version 2.0.0.

Two consequences follow. First, cupy 13.3.0 and numpy 1.24.0 are not floating ranges, so you are committing to a CUDA generation and a numpy ABI that other projects in the same environment will have to match. Second, the list mixes runtime needs with documentation and serving tools. Sphinx, Flask and pywebio are in the same file as the deep learning stack, which makes it harder to reason about a minimal install. The README links to an Installation section near the top of the file.

JIT compilation brings its own debugging cost

The headline capability is JIT compilation, described in the README as just-in-time compilation for dynamic code modification without pre-compilation overhead. The point of a JIT framework is that the model is compiled at the moment it runs, which changes where failures appear. A graph convolution written with a wrong shape or a bad sparse layout is not caught at import time, it fails during the traced execution of the model, and the resulting stack trace refers to generated code rather than to the Python line that caused it.

That cost is amplified by the sparse data path described earlier. The CSC and CSR matrices are built once in a no-grad block and then reused, which is efficient, but it means the graph structure is frozen while the weights are not. Any experiment that changes the edge set after preprocessing has to rebuild those matrices by hand, and the README does not document a helper for that.

The other visible consequence is that the library and the framework are versioned together. The feature list credits the sparse path with CuSparse acceleration, and CuSparse ships with the CUDA toolkit that cupy 13.3.0 targets, so a mismatch between the pinned cupy, the installed CUDA and the Jittor build is the most likely first failure a new user will hit. The README does not publish a troubleshooting or debugging section, and there is no documentation of a CPU-only fallback path, so that check belongs in your evaluation plan rather than in the docs.

Where JittorGeometric is the wrong tool

The clearest case against is a team with an existing PyTorch or TensorFlow codebase. The model API here is Jittor native, with execute methods, jittor.nn layers, and a step driven by optimizer.step(loss). There is no compatibility shim described in the repository, and nothing in the README suggests one. Migrating a research codebase means rewriting the model definitions, not swapping an import.

The second case is a production inference service. The requirements list pulls in Flask 3.1.0 and pywebio 1.8.3, and the repository contains a frontend_demo directory, which are signs of a research and demo orientation rather than a hardened serving stack. The README also states the goal plainly: support for GNN research and applications, with the emphasis on research in the model zoo and the examples.

A third limitation is release hygiene. The repository has no GitHub releases, so there is no tagged artifact to pin, no changelog and no semantic versioning of the distribution beyond the 2.0.0 string in setup.py. The last push was on 2026-06-03, and the branch in use is main. For a library, that means you depend on a commit rather than a release, and the version 2.0 label in the README describes a rewrite whose contents can change under you without a version bump.

How it differs from the graph library most people already have

The natural comparison is with PyTorch Geometric, and the difference is not about model coverage so much as about where the graph lives. PyTorch Geometric is the graph library most people reach for with PyTorch, while JittorGeometric ships its own operator layer: the setup.py package_data includes C++ sources and headers under ops/cpp, and the conversions used in the quick tour, cootocsr and cootocsc, are exposed as first class functions in jittor_geometric.ops. Sparse work goes through those kernels and, per the feature list, CuSparse.

The second difference is the training substrate. JittorGeometric 2.0 lists distributed training over MPI for multi-GPU and multi-node runs, and the examples directory contains a dist_gcn folder to match. A framework that compiles models at runtime has to be careful about how work is partitioned across processes, because each process compiles its own copy of the graph, and that cost is what the MPI path is designed to absorb.

The third difference is the ecosystem it plugs into. The dependency list is Jittor first: jittor at 1.3.9.14, cupy for the GPU path, and an Ascend-GNN for NPU that has no equivalent in a CUDA oriented stack. If your organisation has an Ascend or domestic hardware roadmap, or if it has already standardised on Jittor, that alignment is the whole argument for this library. If not, you are taking on a pinned Jittor and cupy toolchain to reuse a model zoo that exists elsewhere.

Editorial conclusion

JittorGeometric 2.0 fits research groups that already build on Jittor and want the GCN, GAT, GraphSAGE and molecular model set in one place, with 40 plus model implementations covering classic, spectral, dynamic, molecular and self-supervised families, plus a C++ operator layer and MPI training. It does not fit a team that wants a drop-in replacement for PyTorch Geometric, because the model API uses the Jittor execution model with execute methods and the requirements pin jittor to 1.3.9.14 with cupy 13.3.0 and numpy 1.24.0. Before adopting it, check whether the pinned jittor and cupy versions build on your CUDA and Python combination, and read the Installation section of the README, since the repository ships no tagged releases and the last push was on 2026-06-03, so the only artifact to pin is a commit.

Frequently asked questions

Official sources

1. AlgRUC/JittorGeometric on GitHub
2. Issues
3. License: Apache-2.0
4. README