Task 17: Implement ONNX Export and Export-Parity Validation
Implement a clean model-export path for GNN4Colliders, focused first on exporting supported ROOT-GNN inference models to ONNX and validating exported-model parity against native PyTorch inference.
This task builds on:
Task 8: ROOT-GNN model and transfer/fine-tuning
Task 9: classification tasks and output semantics
Task 11: checkpoint loading
Task 12: evaluation and inference
Task 13: Hydra configuration and CLI
Task 14: distributed execution
Task 16: documentation and migration closure
Before making changes, read:
AGENTS.md
README.md
docs/architecture.md
docs/migration.md
Then inspect:
src/gnn4colliders/models/root_gnn/
src/gnn4colliders/tasks/
src/gnn4colliders/training/
src/gnn4colliders/inference/
src/gnn4colliders/cli/
src/gnn4colliders/graphs/
configs/
tests/unit/
tests/integration/
tests/parity/
legacy/root_gnn_dgl/
Also inspect whether the legacy repository contains an active ONNX/export script and characterize only the export behavior actually used.
Do not redesign the ROOT-GNN architecture merely to make ONNX export easier.
Goal
Provide a robust export path:
checkpoint
↓
reconstruct ROOT-GNN model
↓
adapt graph/model inputs into exportable representation
↓
ONNX export
↓
ONNX Runtime inference
↓
compare against PyTorch
Support at minimum:
ROOT-GNN pretrained multiclass model
ROOT-GNN fine-tuned binary model
The export path must preserve inference semantics.
Do not treat successful file creation as sufficient.
Export is only considered successful when ONNX Runtime outputs match native PyTorch outputs within justified tolerances.
1. Add an export package
Prefer a small structure such as:
src/gnn4colliders/export/
__init__.py
onnx.py
adapters.py
or a similarly compact organization.
Possible responsibilities:
onnx.py
export/load/validate ONNX model
adapters.py
convert native ROOT-GNN graph inputs into an export-friendly tensor interface
Do not create a generic deployment framework.
2. Keep export separate from normal model execution
The normal ROOT-GNN model should continue consuming its established graph/batch representation.
Do not rewrite production inference around ONNX constraints.
If an export-specific wrapper is needed, isolate it.
Conceptually:
native EdgeNetwork
↓
ExportAdapter
↓
tensor-only forward signature
↓
ONNX
The adapter exists for export only.
3. Characterize DGL export constraints first
Before writing an adapter, determine whether the current DGL model can be exported directly.
Do not assume DGL graph operations are ONNX-exportable.
Perform a small proof-of-concept.
If direct export fails or produces unsupported operators, document that finding and implement a tensor-based export adapter.
Do not keep a fragile direct-DGL export path merely because a file can be generated.
4. Tensor-based graph representation
If needed, define an explicit export representation.
Conceptually:
node_features
edge_features
edge_src
edge_dst
global_features
graph_index / batch_index
Use the minimum tensors required to reproduce Task 8 message passing.
The export representation must preserve:
node ordering
edge ordering
graph membership
global-feature association
Do not encode hidden positional assumptions beyond those already validated.
5. Export adapter API
Prefer something conceptually like:
adapter = RootGNNExportAdapter(model)
logits = adapter(
node_features,
edge_features,
edge_src,
edge_dst,
graph_index,
global_features,
)
The exact signature may differ.
Keep it explicit and tensor-only where practical.
Do not expose EventMetadata to the exported model.
6. Metadata is not model input
Do not export:
fold
weight
sample_id
source_file
as model inputs unless the model actually consumes them.
Task/event metadata remains outside the ONNX computational graph.
7. Preserve raw-logit output
The exported model should return the same raw logits as native PyTorch.
Do not include:
sigmoid
softmax
thresholding
argmax
inside the ONNX model unless the established production model itself includes them.
Task 9 remains the owner of postprocessing semantics.
8. Export pretrained multiclass model
Support exporting the active multiclass ROOT-GNN model.
Verify:
output shape = [batch_size, num_classes]
Use the active model config from the checkpoint.
Do not hardcode the historical class count into generic export code.
9. Export fine-tuned binary model
Support exporting:
FineTunedEdgeNetwork
Verify:
output shape = [batch_size, 1]
The exported model must include the transferred backbone plus the replacement classifier.
Do not require the source pretraining classifier.
10. Checkpoint integration
Use Task 11 model reconstruction/loading APIs.
Conceptually:
checkpoint
↓
reconstruct model
↓
load weights
↓
model.eval()
↓
export
Do not duplicate checkpoint parsing in export code.
11. Export from model object
Also support exporting an already-constructed model where practical.
Conceptually:
export_root_gnn_onnx(
model=model,
example_batch=batch,
output_path=path,
)
This makes tests easier and avoids coupling export strictly to disk checkpoints.
12. Export from checkpoint convenience path
Provide a higher-level convenience API if useful:
export_checkpoint_to_onnx(
checkpoint_path,
output_path,
example_batch=...,
)
This should reuse the lower-level export path.
Avoid two independent implementations.
13. Example input construction
ONNX export usually requires representative input tensors.
Provide an explicit helper that converts:
GraphBatch
into:
ExportInputs
or equivalent.
Do not make the export code reconstruct graphs from ROOT files.
The export boundary begins from already-prepared graph/batch data.
14. Export input dataclass
If useful, add:
@dataclass(frozen=True)
class RootGNNExportInputs:
node_features: torch.Tensor
edge_features: torch.Tensor
edge_src: torch.Tensor
edge_dst: torch.Tensor
graph_index: torch.Tensor
global_features: torch.Tensor | None
Adjust fields to the actual algorithm.
Keep this internal unless it has real user value.
15. Batching semantics
Support batched graphs if practical.
The adapter must reproduce DGL batching semantics exactly.
Explicitly handle:
which node belongs to which graph
which edge belongs to which graph
per-graph pooling
global features per graph
Do not infer graph boundaries from edge ordering alone if a robust explicit representation exists.
16. Graph membership representation
Use a stable tensor representation for graph membership.
For example:
node_batch_index
edge_batch_index
or offsets such as:
node_splits
edge_splits
Choose the representation that best matches export/runtime compatibility.
Document it.
17. Dynamic graph sizes
Support dynamic:
number of nodes
number of edges
batch size
where the chosen ONNX exporter/runtime supports it reliably.
Use dynamic axes or symbolic dimensions as appropriate.
Do not lock the exported model to one tiny example graph unless unavoidable.
18. Dynamic-axis validation
Test the same ONNX file on at least:
one graph size used for export
a different node/edge count
a different batch size
where supported.
Successful inference on the export example alone is not enough.
19. Edge ordering
Preserve the graph edge order validated in Task 5.
The tensor adapter should derive:
edge_src
edge_dst
edge_features
consistently.
Do not reorder edges during export conversion unless proven numerically irrelevant and explicitly documented.
20. Node aggregation semantics
Reproduce the Task 8 aggregation exactly.
If native code uses DGL sum aggregation, the export adapter must implement equivalent tensor operations.
Potential mechanisms include:
index_add
scatter-like tensor operations
segment reductions
Use operators that export cleanly to ONNX.
Do not change sum aggregation to mean/max.
21. Global pooling semantics
Reproduce native:
mean pooling of node representations
mean pooling of edge representations
global update
or the actual Task 8 behavior exactly.
Pay special attention to varying graph sizes.
Do not compute a mean over the entire batch.
Pooling must remain per graph.
22. Empty graphs
Characterize whether empty-node or empty-edge graphs are supported by the current production pipeline.
If they are not valid inputs, validate and reject them before export/inference.
If they are supported, add explicit tests.
Do not invent unsupported empty-graph semantics solely for ONNX.
23. Single-node graph
Task 5 established that a single-node fully connected graph has:
1 node
0 edges
Ensure the export adapter either supports this correctly or rejects it with a clear documented limitation.
Do not assume every graph has at least one edge.
24. Exportable operations
Prefer common ONNX-supported tensor operations.
Do not introduce custom ONNX ops for the first implementation.
Avoid deployment environments that require custom DGL operators.
The exported graph should be as portable as practical.
25. ONNX opset
Choose an explicit supported ONNX opset.
Do not leave it implicit.
Select the lowest practical opset that supports the required tensor operations and current exporter/runtime.
Document the chosen opset and why.
26. Exporter API
Use the current supported PyTorch ONNX export path available in the project environment.
Do not depend on obsolete exporter APIs if the current PyTorch stack provides a stable newer path.
Keep exporter-specific logic localized in:
export/onnx.py
27. ONNX dependency management
Add ONNX tooling as an optional export/development dependency rather than forcing every GNN4Colliders installation to include it.
Prefer an extra such as:
export
or:
onnx
Use the naming convention that best matches the existing pyproject.toml.
Possible dependencies:
onnx
onnxruntime
Add only what is actually required.
28. CPU ONNX Runtime validation
The automated test suite should validate exported models using CPU ONNX Runtime.
Do not require CUDA ONNX Runtime in normal tests.
CUDA Runtime benchmarking may be optional/manual.
29. ONNX model validation
After export, run ONNX structural validation where supported.
For example:
load model
checker validation
Fail clearly if the graph is invalid.
Do not assume a successfully written file is valid.
30. PyTorch vs ONNX parity
This is the primary correctness requirement.
For the same deterministic inputs:
PyTorch model
->
reference logits
ONNX Runtime
->
exported logits
Compare outputs.
Use appropriate tolerances.
For float32 CPU inference, start with tight tolerances and loosen only if justified by operator differences.
31. Multiclass parity test
Add an export parity test for active multiclass pretraining.
Use:
fixed model weights
deterministic graph inputs
multiple graphs if possible
Verify ONNX logits match PyTorch.
32. Fine-tuning parity test
Add parity for the binary fine-tuned model.
Verify transferred backbone plus new classifier matches native PyTorch.
Test both construction from:
live model
checkpoint
where practical.
33. Multiple processing steps
Ensure parity covers:
processing_steps > 1
so the export graph actually exercises repeated message passing.
Do not validate only a trivial zero/one-step configuration.
34. Global-feature parity
If global features are active, export and test them.
Test:
with global features
and, if supported by the current model:
without global features
Do not invent dummy globals to avoid handling the real interface.
35. Dropout/eval semantics
Export only in evaluation mode.
Ensure:
model.eval()
before export.
Dropout must behave as inference, not training.
Add a repeated native/export parity test if useful.
36. Do not export training
Do not export:
loss
backward
optimizer
scheduler
training loop
The ONNX graph is inference-only.
37. Task postprocessing stays outside ONNX
Do not include Task 9 metrics or weighting.
For a binary exported model:
ONNX -> logits
Python task -> sigmoid / threshold
For multiclass:
ONNX -> logits
Python task -> scores / argmax
This preserves separation of concerns.
38. Optional postprocessing helper
It is acceptable to provide an inference helper that takes ONNX logits and applies the existing Task 9 Python task logic.
Do not create a second independent definition of score semantics.
39. ONNX inference wrapper
Add a lightweight runtime wrapper if useful.
Conceptually:
runner = OnnxPredictor(path)
logits = runner.predict(export_inputs)
Keep it small.
Do not duplicate the full Task 12 Predictor.
40. Runtime wrapper responsibilities
An ONNX runtime helper may own:
loading session
mapping tensor names
NumPy conversion
running inference
returning logits
It should not own:
ROOT reading
graph construction
task metrics
checkpoint handling
41. Input/output names
Use stable human-readable ONNX tensor names.
For example:
node_features
edge_features
edge_src
edge_dst
node_batch
edge_batch
global_features
logits
Use the actual final representation.
Do not expose meaningless generated names when explicit names are easy to provide.
42. Export metadata
Store useful metadata alongside the exported ONNX model.
Possible approaches:
ONNX model metadata
sidecar JSON
At minimum consider:
model family
checkpoint schema version
feature schema version
graph schema version
model config
task type
ONNX opset
GNN4Colliders version
Keep it compact.
43. Sidecar metadata
If ONNX's metadata API is insufficient for structured values, write:
model.onnx
model.onnx.json
or equivalent.
Do not serialize large model tensors twice.
44. Export provenance
If exporting from a checkpoint, record a checkpoint identifier/path or digest as informational provenance.
Do not make absolute local filesystem paths mandatory for model use.
A basename or optional hash may be preferable.
45. Schema compatibility
Export metadata should include:
feature schema version
graph schema version
so future consumers can detect incompatible preprocessing.
Do not expect ONNX itself to construct collider features from raw ROOT branches.
46. Exported-model contract
Document clearly that the exported ROOT-GNN expects already-constructed numerical graph tensors.
The deployment flow is:
raw event
↓
GNN4Colliders preprocessing
↓
export tensor representation
↓
ONNX model
unless a future task exports preprocessing too.
Do not imply model.onnx accepts ROOT files.
47. No preprocessing export
Do not export:
ROOT reading
Awkward transforms
object selection
feature scaling from raw branches
graph construction from raw physics objects
into ONNX in this task.
Keep the scope to the trained neural network computation.
48. CLI export command
Now add:
gnn4colliders export
through the Task 13 CLI.
Keep it thin.
Conceptually:
uv run gnn4colliders export \
export.checkpoint=/path/to/model.pt \
export.output=model.onnx
Use the final Hydra/config conventions.
49. Export config group
Add:
configs/export/
onnx.yaml
or equivalent.
Possible fields:
export:
format: onnx
checkpoint: null
output: model.onnx
opset: ...
Only expose meaningful stable options.
50. Example input for CLI export
The CLI needs a representative graph/batch to establish export shapes.
Choose a clean explicit approach.
Possibilities include:
use a configured dataset and first batch
use a saved export-input fixture
accept dimensions plus a generated valid dummy graph
Prefer using the actual configured preprocessing/data pipeline when practical, because it validates the real input representation.
Do not silently export using a tiny shape that cannot generalize.
51. export command responsibility
The CLI should:
resolve config
load checkpoint
construct/reuse example batch
convert to export inputs
export ONNX
validate ONNX
optionally run parity check
write metadata
Actual export math stays in gnn4colliders.export.
52. CLI parity validation
By default, after exporting, run at least one PyTorch-vs-ONNX comparison unless this would be prohibitively expensive.
Prefer failing export if parity validation fails.
If a skip option exists, make it explicit.
Do not silently produce an unvalidated export.
53. Output overwrite behavior
Do not overwrite an existing ONNX file unexpectedly.
Use the project’s existing output policy.
Either:
refuse
require explicit overwrite
write unique path
Keep this non-interactive for batch jobs.
54. Export on CPU
The baseline export workflow should work on CPU.
A user should not need a GPU merely to convert a checkpoint to ONNX.
Map the model/example tensors appropriately.
55. Export GPU-trained checkpoints on CPU
Checkpoints created on CUDA should be exportable on a CPU machine using Task 11 map_location behavior where architecture dependencies permit.
Add a test where practical.
56. ONNX Runtime device independence
Automated parity should use CPU ONNX Runtime.
Do not assume a CUDA execution provider.
Document optional CUDA runtime use separately only if validated.
57. Dynamic shapes vs portability
Prefer an export that supports realistic variable graph sizes.
However, if specific ONNX operators/exporter limitations prevent fully dynamic shapes, document exact constraints.
Do not pretend a fixed-shape model is dynamic.
If necessary, support a clearly named fixed-shape mode only as a fallback.
58. Avoid padding redesign
Do not reintroduce legacy fixed:
16000 nodes
104000 edges
padding merely to simplify export.
Use the clean new graph representation.
Only use fixed padding if a real downstream deployment target requires it and this is explicitly documented.
59. ONNX graph inspection
Inspect the exported graph.
Verify it does not accidentally include:
training-only branches
constant example-specific graph topology
hardcoded example batch size
hardcoded example node count
unless those dimensions are intentionally fixed.
Add tests for dynamic behavior rather than relying only on manual inspection.
60. Operator compatibility
Record the major ONNX operators needed by the exported model.
Do not manually optimize the ONNX graph in this task.
If an operator is poorly supported by the target runtime, adjust the adapter using equivalent standard tensor operations where possible.
61. Avoid custom scatter dependencies if possible
If aggregation requires scatter behavior, prefer ONNX-exportable native PyTorch operations.
Do not add PyTorch Scatter solely for export unless necessary.
Keep runtime dependencies minimal.
62. Performance is secondary to parity
Benchmark ONNX inference if useful, but do not optimize the adapter solely for speed in this task.
The primary goal is:
portable export + correctness
not beating native DGL inference.
63. Optional ONNX benchmark
If straightforward, add a small benchmark:
benchmarks/benchmark_onnx.py
Compare:
PyTorch CPU inference
ONNX Runtime CPU inference
for representative inputs.
Do not make performance claims without measurements.
64. Unit tests
Add focused tests under:
tests/unit/export/
Suggested coverage:
test_export_inputs.py
test_export_adapter.py
test_onnx_export.py
test_onnx_runtime.py
Use fewer files if clearer.
65. Export-input conversion tests
Verify conversion from GraphBatch preserves:
node features
edge features
edge src/dst
graph membership
global features
graph count
Compare directly against the native DGL graph.
66. Adapter native parity
Before ONNX export, compare:
native DGL EdgeNetwork
against:
tensor-only ExportAdapter
using the same weights and graph.
This isolates adapter correctness from ONNX exporter issues.
Require tight parity.
67. Adapter multiclass test
Test multiclass native-vs-adapter parity.
Use multiple processing steps and multiple graphs.
68. Adapter fine-tuning test
Test fine-tuned binary native-vs-adapter parity.
Verify the new classifier is represented correctly.
69. ONNX file smoke test
Export a tiny valid model.
Verify:
file exists
model loads
ONNX checker succeeds
runtime session initializes
Do not stop there; run numerical parity too.
70. ONNX numerical parity test
Compare native PyTorch and ONNX Runtime logits.
Use deterministic fixed inputs.
Test at least:
one graph
multiple graphs
different graph size
where dynamic export supports them.
71. Batch-size dynamic test
If batch dimension is declared dynamic, test:
batch size 1
batch size > 1
with the same exported model.
72. Node-count dynamic test
If node count is dynamic, test a different number of nodes from the export example.
Because edges scale as:
N * (N - 1)
this also validates dynamic edge count.
73. Single-node test
Where supported, validate:
N = 1
E = 0
through:
native model
adapter
ONNX Runtime
If ONNX cannot support this case cleanly, validate the rejection path and document it.
74. Checkpoint export integration test
Add:
model
->
Task 11 checkpoint
->
fresh export load
->
ONNX
->
runtime logits
Compare against the original model.
This proves the user-facing workflow, not just direct model export.
75. Fine-tuned checkpoint export integration
Explicitly test:
multiclass pretrained backbone
->
fine-tuned binary checkpoint
->
export
->
ONNX binary logits
This is a required real workflow.
76. CLI export integration test
Exercise the final CLI with temporary files.
Conceptually:
checkpoint
->
gnn4colliders export
->
model.onnx
->
ONNX Runtime validation
Keep the model/data tiny.
Do not require GPU.
77. Metadata sidecar test
If metadata is written, verify:
model family
model/task config
schema versions
opset
are present and parseable.
Do not test ephemeral timestamp strings too strictly.
78. Invalid checkpoint behavior
Test:
wrong model family
unsupported checkpoint schema
missing model state
incompatible ROOT-GNN config
Raise clear export errors.
Do not emit partial ONNX files after validation failure.
79. Atomic export write
Where practical, export to a temporary file and only move into the final path after:
export succeeds
ONNX validation succeeds
parity validation succeeds
Avoid leaving a seemingly valid final artifact after failure.
80. Cleanup temporary files
Ensure failed export attempts clean up temporary ONNX/metadata artifacts where practical.
Do not leave test litter in output directories.
81. Documentation
Add:
docs/export.md
or an equivalent clear section.
Document:
supported models
export command
tensor input contract
dynamic dimensions
raw-logit output
metadata sidecar
ONNX Runtime validation
known limitations
82. README export section
Add a concise example:
uv run gnn4colliders export \
export.checkpoint=/path/to/model.pt \
export.output=model.onnx
Use the actual final syntax.
Link to detailed export documentation.
83. Architecture documentation
Update:
docs/architecture.md
with the export boundary:
GraphBatch
->
ROOT-GNN ExportAdapter
->
tensor-only representation
->
ONNX model
Make clear this is an inference/deployment adapter, not the primary graph representation.
84. Migration documentation
Update:
docs/migration.md
to mark ONNX/export status accurately.
If legacy export behavior cannot be reproduced exactly, describe the new supported contract.
Do not mark legacy deployment parity complete unless validated.
85. Legacy export parity
If the legacy repository has an active ONNX export path, characterize:
input representation
output semantics
fixed/dynamic shapes
supported model type
Compare only externally relevant behavior.
Do not preserve awkward legacy export implementation solely for implementation parity.
86. No TensorRT in this task
Do not implement:
TensorRT
Torch-TensorRT
OpenVINO
This task is ONNX only.
A portable ONNX artifact is the foundation for later deployment backends.
87. No preprocessing deployment framework
Do not build C++/CUDA ROOT preprocessing for ONNX deployment.
The exported model contract begins with processed graph tensors.
Keep scope bounded.
88. No quantization
Do not add:
INT8 quantization
dynamic quantization
QAT
in this task.
Quantization changes numerical behavior and deserves its own dedicated task if needed.
89. No model simplification dependency
Do not add ONNX graph-simplifier tooling as a required dependency.
Only use it if a proven exporter issue requires it.
Prefer standard ONNX output first.
90. No distributed export
Only rank 0 should perform export if the CLI is accidentally launched under a distributed environment.
Do not export one ONNX file per DDP rank.
Export operates on the normalized underlying model.
91. DDP checkpoint export
Explicitly verify a checkpoint produced during DDP training can be exported.
Task 14 should already normalize state dicts.
Add an integration test or reuse an existing normalized DDP checkpoint fixture.
92. Export config validation
Validate:
checkpoint exists
format is supported
output extension matches ONNX
opset is supported
model family supports export
required example/input source is available
Fail before expensive work where possible.
93. Public API
Expose a small intended API.
For example:
from gnn4colliders.export import (
export_root_gnn_onnx,
RootGNNExportAdapter,
)
Only expose the adapter if users genuinely need it.
Keep internal tensor conversion helpers private where practical.
94. Optional extra installation docs
If ONNX dependencies are optional, document:
uv sync --extra root-gnn --extra onnx
or the actual selected extra name.
Do not require ONNX dependencies for users who only train ROOT-GNN.
95. Import hygiene
Installing only the normal ROOT-GNN extra should not make:
import gnn4colliders
fail because ONNX Runtime is absent.
Keep optional dependency imports inside export functionality.
Raise a clear actionable error if export is requested without the required extra.
96. Optional dependency error
Prefer an error conceptually like:
ONNX export requires the 'onnx' optional dependency.
Install with: uv sync --extra root-gnn --extra onnx
Use the actual package-extra name.
Do not emit a cryptic ModuleNotFoundError where a clear project error is easy.
97. Validate no native behavior regression
After implementing export, run the full existing test suite.
ONNX support must not change:
training
checkpoint loading
fine-tuning
inference
DDP
The native path remains authoritative.
98. Validation
Run focused export tests:
uv run pytest tests/unit/export -v
Run export integration tests:
uv run pytest tests/integration -k "export or onnx" -v
Run model parity tests:
uv run pytest tests/parity -v
Run the full suite:
uv run pytest
Run lint/format checks:
uv run ruff check src tests
uv run ruff format --check src tests
Validate CLI:
uv run gnn4colliders export --help
Perform a manual tiny export:
uv run gnn4colliders export \
export.checkpoint=/path/to/test/checkpoint.pt \
export.output=/tmp/root_gnn.onnx
Then validate the generated ONNX model with the implemented runtime/parity utility.
Inspect:
git status
git diff
Verify:
- native ROOT-GNN model mathematics remain unchanged
- no task/loss semantics moved inside ONNX
- no metadata/tracking inputs were added to the model
- no TensorRT/quantization/custom-op scope was introduced
- ONNX dependencies remain optional
- no large generated ONNX files are tracked
- no unrelated changes are included
Completion criteria
Task 17 is complete when:
- An isolated export package exists.
- The native ROOT-GNN model remains unchanged in its public inference behavior.
- A tensor-only export adapter exists if DGL cannot be exported directly.
- Adapter outputs match native
EdgeNetworkoutputs. - Adapter outputs match native fine-tuned model outputs.
- Multiclass ROOT-GNN can be exported.
- Binary fine-tuned ROOT-GNN can be exported.
- Raw logits are the exported model output.
- Graph batching semantics are preserved.
- Node/edge/global message-passing semantics are preserved.
- Variable graph size is supported to the documented extent.
- Variable batch size is supported to the documented extent.
- ONNX opset is explicit.
- ONNX structural validation succeeds.
- ONNX Runtime CPU inference works.
- PyTorch vs ONNX multiclass parity passes.
- PyTorch vs ONNX binary fine-tuning parity passes.
- Checkpoint-based export works.
- DDP-produced normalized checkpoints can be exported.
- Export metadata/schema information is written.
gnn4colliders exportworks through the stable Python export API.- Invalid/incompatible exports fail clearly.
- ONNX dependencies are optional.
- Normal package imports work without ONNX dependencies.
- Export documentation exists.
- README contains a validated export example.
- Full existing tests still pass.
- Legacy source remains untouched.
Completion report
Report:
- files created
- files modified
- direct-DGL export findings
- export adapter design
- tensor input contract
- graph-membership representation
- dynamic dimension support
- supported ROOT-GNN model types
- ONNX opset
- export API
- checkpoint export API
- CLI export workflow
- ONNX metadata/sidecar format
- native-vs-adapter parity results
- adapter-vs-ONNX parity results
- multiclass export results
- binary fine-tuning export results
- single-node/zero-edge behavior
- variable node/edge test results
- variable batch-size test results
- DDP checkpoint export results
- ONNX Runtime validation results
- numerical tolerances used
- optional dependency changes
- known export limitations
- behavior intentionally deferred
- validation commands and results
After validation succeeds, create one Git commit containing only Task 17 changes.
Use:
feat: add ROOT-GNN ONNX export
Before committing, inspect the final diff and ensure no unrelated files or generated ONNX artifacts are included.