# 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:

```text id="70ycbf"
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:

```text id="qt0f76"
AGENTS.md
README.md
docs/architecture.md
docs/migration.md
```

Then inspect:

```text id="g7fohc"
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:

```text id="ww6x60"
checkpoint
    ↓
reconstruct ROOT-GNN model
    ↓
adapt graph/model inputs into exportable representation
    ↓
ONNX export
    ↓
ONNX Runtime inference
    ↓
compare against PyTorch
```

Support at minimum:

```text id="de5e55"
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:

```text id="f5m7kt"
src/gnn4colliders/export/
    __init__.py
    onnx.py
    adapters.py
```

or a similarly compact organization.

Possible responsibilities:

```text id="re1n1w"
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:

```text id="pnbm0r"
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:

```text id="2dqa9s"
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:

```text id="we9d1w"
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:

```python id="kz3uy0"
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:

```text id="dbv0zk"
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:

```text id="oip3wc"
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:

```text id="kkvpgz"
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:

```text id="cr9zyu"
FineTunedEdgeNetwork
```

Verify:

```text id="qq26fr"
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:

```text id="15bqxe"
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:

```python id="d7kmlu"
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:

```python id="j37f3n"
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:

```text id="pkyvh2"
GraphBatch
```

into:

```text id="1bg1fm"
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:

```python id="rocyja"
@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:

```text id="i3bbf1"
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:

```text id="hi8tqm"
node_batch_index
edge_batch_index
```

or offsets such as:

```text id="s3b23o"
node_splits
edge_splits
```

Choose the representation that best matches export/runtime compatibility.

Document it.

---

# 17. Dynamic graph sizes

Support dynamic:

```text id="w3f4xn"
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:

```text id="3g043o"
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:

```text id="9obebd"
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:

```text id="akrzpu"
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:

```text id="65w0yl"
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:

```text id="wsdhnc"
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:

```text id="lb6wc8"
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:

```text id="deu91j"
export
```

or:

```text id="5em87q"
onnx
```

Use the naming convention that best matches the existing `pyproject.toml`.

Possible dependencies:

```text id="e8gnwn"
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:

```text id="34m6h9"
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:

```text id="jx7lus"
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:

```text id="eclkeo"
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:

```text id="bo2ura"
live model
checkpoint
```

where practical.

---

# 33. Multiple processing steps

Ensure parity covers:

```text id="06jwtq"
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:

```text id="mj6a6v"
with global features
```

and, if supported by the current model:

```text id="uv9157"
without global features
```

Do not invent dummy globals to avoid handling the real interface.

---

# 35. Dropout/eval semantics

Export only in evaluation mode.

Ensure:

```python id="rhp02x"
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:

```text id="jzct8f"
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:

```text id="z73v2h"
ONNX -> logits
Python task -> sigmoid / threshold
```

For multiclass:

```text id="yyqjek"
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:

```python id="zis9l9"
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:

```text id="9r6aqc"
loading session
mapping tensor names
NumPy conversion
running inference
returning logits
```

It should not own:

```text id="0fjm8y"
ROOT reading
graph construction
task metrics
checkpoint handling
```

---

# 41. Input/output names

Use stable human-readable ONNX tensor names.

For example:

```text id="yaw4v1"
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:

```text id="btxb1f"
ONNX model metadata
sidecar JSON
```

At minimum consider:

```text id="lmndxr"
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:

```text id="2j9vch"
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:

```text id="btmtr1"
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:

```text id="tfxfjz"
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:

```text id="s6l78d"
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:

```bash id="fs047y"
gnn4colliders export
```

through the Task 13 CLI.

Keep it thin.

Conceptually:

```bash id="zqlqrr"
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:

```text id="8e3w0v"
configs/export/
    onnx.yaml
```

or equivalent.

Possible fields:

```yaml id="o9pmef"
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:

```text id="ef7hnb"
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:

```text id="u7ykk2"
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:

```text id="h8k3q5"
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:

```text id="ygwyq3"
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:

```text id="bmbynj"
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:

```text id="ph78rw"
portable export + correctness
```

not beating native DGL inference.

---

# 63. Optional ONNX benchmark

If straightforward, add a small benchmark:

```text id="8qp54e"
benchmarks/benchmark_onnx.py
```

Compare:

```text id="ucg05t"
PyTorch CPU inference
ONNX Runtime CPU inference
```

for representative inputs.

Do not make performance claims without measurements.

---

# 64. Unit tests

Add focused tests under:

```text id="bu7qgo"
tests/unit/export/
```

Suggested coverage:

```text id="sjk6ld"
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:

```text id="7oxvoc"
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:

```text id="r946qe"
native DGL EdgeNetwork
```

against:

```text id="13amoi"
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:

```text id="yd7du1"
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:

```text id="4umw7v"
one graph
multiple graphs
different graph size
```

where dynamic export supports them.

---

# 71. Batch-size dynamic test

If batch dimension is declared dynamic, test:

```text id="ws79d7"
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:

```text id="46vfip"
N * (N - 1)
```

this also validates dynamic edge count.

---

# 73. Single-node test

Where supported, validate:

```text id="l6zhtf"
N = 1
E = 0
```

through:

```text id="b5vlcx"
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:

```text id="jvc91c"
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:

```text id="4zjsai"
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:

```text id="z95kom"
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:

```text id="a8l6zc"
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:

```text id="qxah3i"
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:

```text id="mk0s8v"
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:

```text id="8xjvnm"
docs/export.md
```

or an equivalent clear section.

Document:

```text id="xj73lt"
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:

```bash id="fcxzsb"
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:

```text id="qahvah"
docs/architecture.md
```

with the export boundary:

```text id="uydv0k"
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:

```text id="2wglae"
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:

```text id="sb3fad"
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:

```text id="xd00dm"
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:

```text id="12o86l"
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:

```text id="onq10b"
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:

```python id="2e4a4t"
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:

```bash id="5iiijm"
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:

```python id="z1mofz"
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:

```text id="20w0as"
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:

```text id="qr6203"
training
checkpoint loading
fine-tuning
inference
DDP
```

The native path remains authoritative.

---

# 98. Validation

Run focused export tests:

```bash id="wd1h8a"
uv run pytest tests/unit/export -v
```

Run export integration tests:

```bash id="93g5qh"
uv run pytest tests/integration -k "export or onnx" -v
```

Run model parity tests:

```bash id="p5c8fy"
uv run pytest tests/parity -v
```

Run the full suite:

```bash id="lix9ck"
uv run pytest
```

Run lint/format checks:

```bash id="1o5l3y"
uv run ruff check src tests
uv run ruff format --check src tests
```

Validate CLI:

```bash id="0qr5zt"
uv run gnn4colliders export --help
```

Perform a manual tiny export:

```bash id="g1ur2s"
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:

```bash id="fskv6x"
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:

```text id="6p7aqi"
feat: add ROOT-GNN ONNX export
```

Before committing, inspect the final diff and ensure no unrelated files or generated ONNX artifacts are included.