[tmva][sofie] Evolve the SOFIE GNN support from graph_nets to PyTorch + ONNX - #23144
Open
guitargeek wants to merge 2 commits into
Open
[tmva][sofie] Evolve the SOFIE GNN support from graph_nets to PyTorch + ONNX#23144guitargeek wants to merge 2 commits into
guitargeek wants to merge 2 commits into
Conversation
guitargeek
requested review from
bellenot,
couet,
dpiparo,
lmoneta and
vepadulano
as code owners
August 25, 2026 08:54
Test Results 23 files 23 suites 3d 18h 14m 41s ⏱️ For more details on these failures, see this check. Results for commit 0ef139e. ♻️ This comment has been updated with latest results. |
guitargeek
force-pushed
the
sofie-gnn-onnx-migration
branch
2 times, most recently
from
August 31, 2026 11:56
a478b25 to
cd477cb
Compare
couet
removed their request for review
September 1, 2026 09:00
guitargeek
force-pushed
the
sofie-gnn-onnx-migration
branch
from
September 3, 2026 15:58
cd477cb to
f599507
Compare
The SOFIE GNN tutorials defined their model with DeepMind graph_nets and dm-sonnet, whose last releases date from 2020 and which cannot be installed together with current Python and TensorFlow versions anymore. Define the same model with plain PyTorch instead: an Encode-Process-Decode graph network with the identical architecture (graph-network blocks following Battaglia et al., arXiv:1806.01261, with 4-layer ReLU MLP update functions, LayerNorm in the core network and summation aggregation). The components are exported to ONNX with the torch.export-based exporter and the inference code is generated with the SOFIE ONNX parser, so the GNN-specific SOFIE classes and the graph_nets parser are no longer involved. The graph-network block maps onto standard ONNX operators: Gemm+Relu, Concat, Gather of the sender/receiver node features, ScatterElements(reduction=add) for the edge aggregation (the equivalent of unsorted_segment_sum), ReduceSum and LayerNormalization. - TMVA_SOFIE_GNN.py validates the generated inference against PyTorch (agreement at float32 precision) and compares the execution times. - TMVA_SOFIE_GNN_Parser.py exports the model with dynamic node and edge counts, generates the inference code and writes variable-sized input graphs plus PyTorch reference results to a ROOT file. - TMVA_SOFIE_GNN_Application.C evaluates the generated code in C++ on that file, like before, now with the ONNX-generated session interface. The tutorials are now gated on the torch and onnx Python modules instead of sonnet and graph_nets. 🤖 Done with the help of AI
The SOFIE GNN classes could only be filled by parsing models built with DeepMind graph_nets and dm-sonnet, whose last releases date from 2020: they cannot be installed together with current Python and TensorFlow versions anymore (the corresponding Python test had already silently disappeared from CI, since it is gated on the availability of those modules), and the parser relied on private implementation details of both packages. Equivalent graph-network models defined with PyTorch can be deployed through the SOFIE ONNX parser, whose operator support covers the graph-network building blocks; the migrated TMVA_SOFIE_GNN tutorials demonstrate that workflow and produce inference results identical to this implementation. Removed: - RModel_GNN, RModel_GraphIndependent and the RFunction classes (RFunction_MLP, RFunction_Sum, RFunction_Mean), together with their GNN_Init/GraphIndependent_Init input structs and the EmitGNN / EmitGraphIndependent tests - the GNN parts of RModel_Base (Options::kGNN / kGNNComponent, the GraphType/FunctionTarget/FunctionReducer/FunctionRelation enums and the RModel_GNNBase class) and the corresponding code paths and the weight-file read position in RModel::Generate - the GNN_Data struct with its Copy/Concatenate helpers in SOFIE_common and its emission in the generated-code helpers - the ParseFromMemory Python interface (_gnn.py pythonization) and its test - the dm-sonnet and graph_nets entries in requirements.txt 🤖 Done with the help of AI
guitargeek
force-pushed
the
sofie-gnn-onnx-migration
branch
from
September 4, 2026 17:35
f599507 to
0ef139e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The SOFIE GNN support that Sanjiban Sengupta and Lorenzo Moneta developed made SOFIE one of the first inference engines to generate standalone C++ code for graph networks, built around DeepMind's graph_nets formalism (Battaglia et al., arXiv:1806.01261). Since then the ML ecosystem has moved: graph_nets and dm-sonnet have been unmaintained since 2020, while SOFIE's ONNX operator coverage has grown to the point where it covers the
graph-network building blocks (
Gather,ScatterElementswith add reduction,ReduceSum,LayerNormalization, ...), largely thanks to the operator work driven by that same GNN effort.This PR completes that evolution: the same Encode-Process-Decode graph networks are now defined in PyTorch, exported to ONNX (with dynamic node/edge counts), and deployed through SOFIE's general ONNX pipeline, making the dedicated GNN code path and its dependency on the discontinued packages unnecessary. The migrated
TMVA_SOFIE_GNNtutorials produce inference results bitwise identical to the previous implementation.