GNN4Colliders / tasks /task17.md
ho22joshua's picture
rewriting codebase (#7)
916755e
|
Raw
History Blame
33.5 kB

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:

  1. An isolated export package exists.
  2. The native ROOT-GNN model remains unchanged in its public inference behavior.
  3. A tensor-only export adapter exists if DGL cannot be exported directly.
  4. Adapter outputs match native EdgeNetwork outputs.
  5. Adapter outputs match native fine-tuned model outputs.
  6. Multiclass ROOT-GNN can be exported.
  7. Binary fine-tuned ROOT-GNN can be exported.
  8. Raw logits are the exported model output.
  9. Graph batching semantics are preserved.
  10. Node/edge/global message-passing semantics are preserved.
  11. Variable graph size is supported to the documented extent.
  12. Variable batch size is supported to the documented extent.
  13. ONNX opset is explicit.
  14. ONNX structural validation succeeds.
  15. ONNX Runtime CPU inference works.
  16. PyTorch vs ONNX multiclass parity passes.
  17. PyTorch vs ONNX binary fine-tuning parity passes.
  18. Checkpoint-based export works.
  19. DDP-produced normalized checkpoints can be exported.
  20. Export metadata/schema information is written.
  21. gnn4colliders export works through the stable Python export API.
  22. Invalid/incompatible exports fail clearly.
  23. ONNX dependencies are optional.
  24. Normal package imports work without ONNX dependencies.
  25. Export documentation exists.
  26. README contains a validated export example.
  27. Full existing tests still pass.
  28. Legacy source remains untouched.

Completion report

Report:

  1. files created
  2. files modified
  3. direct-DGL export findings
  4. export adapter design
  5. tensor input contract
  6. graph-membership representation
  7. dynamic dimension support
  8. supported ROOT-GNN model types
  9. ONNX opset
  10. export API
  11. checkpoint export API
  12. CLI export workflow
  13. ONNX metadata/sidecar format
  14. native-vs-adapter parity results
  15. adapter-vs-ONNX parity results
  16. multiclass export results
  17. binary fine-tuning export results
  18. single-node/zero-edge behavior
  19. variable node/edge test results
  20. variable batch-size test results
  21. DDP checkpoint export results
  22. ONNX Runtime validation results
  23. numerical tolerances used
  24. optional dependency changes
  25. known export limitations
  26. behavior intentionally deferred
  27. 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.