feat: add inference and ONNX export adapters

docs/export.md

@@ -0,0 +1,36 @@
+# ROOT-GNN ONNX export
+
+ONNX export is an inference-only adapter for prepared ROOT-GNN graph batches.
+It does not read ROOT files or move feature construction into the deployment
+model. Install the optional dependencies with:
+
+```bash
+uv sync --extra root-gnn --extra onnx
+```
+
+Export a checkpoint with the project CLI:
+
+```bash
+uv run gnn4colliders export \
+  export.checkpoint=/path/to/checkpoint.pt \
+  export.output=model.onnx \
+  data.cache.path=/path/to/graph-cache.pt
+```
+
+The exported model accepts six tensor inputs: `node_features`,
+`edge_features`, `edge_src`, `edge_dst`, `node_batch`, and `global_features`.
+`node_batch` identifies the graph for each node; edge membership is derived
+from `node_batch[edge_src]`. Graph, node, edge, and batch dimensions are
+dynamic. The model returns raw `logits`; task postprocessing and metadata stay
+outside ONNX.
+
+The exporter uses ONNX opset 17, validates the structure with `onnx.checker`,
+and writes compact provenance/schema metadata to `model.onnx.json`. It
+supports `EdgeNetwork` multiclass models and `FineTunedEdgeNetwork` binary
+models. Empty graphs are invalid; single-node graphs (zero edges) are
+represented and pooled with a zero edge contribution.
+
+Direct DGL export was not retained: the active model uses DGL graph mutation
+and reductions (`apply_edges`, `update_all`, and graph pooling) that are not a
+portable ONNX contract. The adapter expresses those operations with standard
+tensor indexing, `index_add`, and per-graph means.

src/gnn4colliders/export/__init__.py

@@ -0,0 +1,21 @@
+"""Optional ONNX export for already-prepared ROOT-GNN graph tensors."""
+
+from .onnx import (
+    RootGNNExportAdapter,
+    RootGNNExportInputs,
+    export_checkpoint_to_onnx,
+    export_root_gnn_onnx,
+    inputs_from_graph_batch,
+    validate_onnx,
+    validate_onnx_runtime,
+)
+
+__all__ = [
+    "RootGNNExportAdapter",
+    "RootGNNExportInputs",
+    "export_checkpoint_to_onnx",
+    "export_root_gnn_onnx",
+    "inputs_from_graph_batch",
+    "validate_onnx",
+    "validate_onnx_runtime",
+]

src/gnn4colliders/export/onnx.py

@@ -0,0 +1,375 @@
+"""Tensor-only ROOT-GNN export and ONNX Runtime validation.
+
+The native model intentionally remains DGL-based.  This module is an
+inference/deployment adapter whose public tensor contract is:
+``node_features``, ``edge_features``, ``edge_src``, ``edge_dst``,
+``node_batch``, and ``global_features``.  ``global_features`` is a required
+runtime input for a stable ONNX signature; for models without configured
+globals it is ignored and may contain one zero column per graph.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Mapping
+
+import torch
+from torch import nn
+
+from gnn4colliders.data.metadata import FEATURE_SCHEMA_VERSION, GRAPH_SCHEMA_VERSION
+from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
+from gnn4colliders.training import CheckpointManager, load_model_weights
+
+ONNX_OPSET = 17
+_INPUT_NAMES = (
+    "node_features",
+    "edge_features",
+    "edge_src",
+    "edge_dst",
+    "node_batch",
+    "global_features",
+)
+
+
+@dataclass(frozen=True)
+class RootGNNExportInputs:
+    """Numerical graph representation consumed by the exported model."""
+
+    node_features: torch.Tensor
+    edge_features: torch.Tensor
+    edge_src: torch.Tensor
+    edge_dst: torch.Tensor
+    node_batch: torch.Tensor
+    global_features: torch.Tensor
+
+    @property
+    def graph_count(self) -> int:
+        return int(self.global_features.shape[0])
+
+    def as_tuple(self) -> tuple[torch.Tensor, ...]:
+        return tuple(getattr(self, name) for name in _INPUT_NAMES)
+
+
+def inputs_from_graph_batch(batch: Any) -> RootGNNExportInputs:
+    """Convert a ``GraphBatch`` to explicit, metadata-free export inputs."""
+    graph = batch.graph if hasattr(batch, "graph") else batch
+    if not hasattr(graph, "ndata") or not hasattr(graph, "edata"):
+        raise TypeError("example batch must contain a DGL graph")
+    node_features = graph.ndata["features"]
+    edge_features = graph.edata["features"]
+    if node_features.ndim != 2 or edge_features.ndim != 2:
+        raise ValueError("graph node and edge features must be rank-2 tensors")
+    try:
+        node_counts = graph.batch_num_nodes().to(dtype=torch.long)
+        edge_counts = graph.batch_num_edges().to(dtype=torch.long)
+        edge_src, edge_dst = graph.edges(order="eid")
+    except AttributeError as error:
+        raise ValueError("export requires a batched DGL graph") from error
+    if (
+        node_counts.numel() == 0
+        or int(node_counts.sum()) == 0
+        or bool((node_counts == 0).any())
+    ):
+        raise ValueError("empty-node graphs are not valid ROOT-GNN export inputs")
+    node_batch = torch.repeat_interleave(
+        torch.arange(len(node_counts), device=node_counts.device), node_counts
+    )
+    edge_batch = torch.repeat_interleave(
+        torch.arange(len(edge_counts), device=edge_counts.device), edge_counts
+    )
+    # The edge batch is deliberately derived from source node membership in the
+    # adapter; checking it here catches malformed graph batches early.
+    if edge_batch.numel() and not torch.equal(edge_batch, node_batch[edge_src]):
+        raise ValueError("DGL edge ordering does not agree with graph membership")
+    globals_ = getattr(batch, "global_features", None)
+    if globals_ is None:
+        globals_ = node_counts[:, None].to(dtype=node_features.dtype)
+    elif globals_.ndim == 1:
+        globals_ = globals_.unsqueeze(0)
+    if globals_.shape[0] != len(node_counts):
+        raise ValueError("global_features must have one row per graph")
+    return RootGNNExportInputs(
+        node_features=node_features,
+        edge_features=edge_features,
+        edge_src=edge_src.to(dtype=torch.long),
+        edge_dst=edge_dst.to(dtype=torch.long),
+        node_batch=node_batch.to(dtype=torch.long),
+        global_features=globals_,
+    )
+
+
+def _mean_by_group(
+    values: torch.Tensor, groups: torch.Tensor, count: int
+) -> torch.Tensor:
+    result = values.new_zeros((count, values.shape[1]))
+    if values.shape[0]:
+        result = result.index_add(0, groups, values)
+        sizes = values.new_zeros((count, 1)).index_add(
+            0, groups, values.new_ones((values.shape[0], 1))
+        )
+        result = result / sizes.clamp_min(1)
+    return result
+
+
+def _sum_by_group(
+    values: torch.Tensor, groups: torch.Tensor, count: int
+) -> torch.Tensor:
+    result = values.new_zeros((count, values.shape[1]))
+    if values.shape[0]:
+        result = result.index_add(0, groups, values)
+    return result
+
+
+class RootGNNExportAdapter(nn.Module):
+    """Reproduce an ``EdgeNetwork``/``FineTunedEdgeNetwork`` with tensors only."""
+
+    def __init__(self, model: nn.Module) -> None:
+        super().__init__()
+        if isinstance(model, FineTunedEdgeNetwork):
+            self.backbone = model.backbone
+            self.classifier = model.classifier
+        elif isinstance(model, EdgeNetwork):
+            self.backbone = model
+            self.classifier = model.classifier
+        else:
+            raise TypeError("model must be EdgeNetwork or FineTunedEdgeNetwork")
+        self.has_global = self.backbone.has_global
+
+    def forward(
+        self,
+        node_features: torch.Tensor,
+        edge_features: torch.Tensor,
+        edge_src: torch.Tensor,
+        edge_dst: torch.Tensor,
+        node_batch: torch.Tensor,
+        global_features: torch.Tensor,
+    ) -> torch.Tensor:
+        graph_count = global_features.shape[0]
+        node_h = self.backbone.node_encoder(node_features)
+        edge_h = self.backbone.edge_encoder(edge_features)
+        if self.has_global:
+            global_h = self.backbone.global_encoder(global_features)
+        else:
+            counts = _sum_by_group(
+                node_features.new_ones((node_features.shape[0], 1)),
+                node_batch,
+                graph_count,
+            )
+            global_h = self.backbone.global_encoder(counts)
+        edge_batch = node_batch[edge_src]
+        for _ in range(self.backbone.n_proc_steps):
+            edge_h = self.backbone.edge_update(
+                torch.cat(
+                    (edge_h, node_h[edge_src], node_h[edge_dst], global_h[edge_batch]),
+                    1,
+                )
+            )
+            node_messages = node_h.new_zeros(node_h.shape).index_add(
+                0, edge_dst, edge_h
+            )
+            node_h = self.backbone.node_update(
+                torch.cat((node_h, node_messages, global_h[node_batch]), 1)
+            )
+            global_h = self.backbone.global_update(
+                torch.cat(
+                    (
+                        global_h,
+                        _mean_by_group(node_h, node_batch, graph_count),
+                        _mean_by_group(edge_h, edge_batch, graph_count),
+                    ),
+                    1,
+                )
+            )
+        return self.classifier(self.backbone.global_decoder(global_h))
+
+
+def _optional_onnx() -> Any:
+    try:
+        import onnx
+    except ImportError as error:
+        raise ImportError(
+            "ONNX export requires the 'onnx' extra. Install with: "
+            "uv sync --extra root-gnn --extra onnx"
+        ) from error
+    return onnx
+
+
+def validate_onnx(path: str | os.PathLike[str]) -> None:
+    """Run ONNX structural validation and fail clearly on invalid artifacts."""
+    onnx = _optional_onnx()
+    model = onnx.load(str(path))
+    onnx.checker.check_model(model)
+
+
+def validate_onnx_runtime(
+    path: str | os.PathLike[str],
+    inputs: RootGNNExportInputs,
+    reference: torch.Tensor,
+    *,
+    rtol: float = 1e-4,
+    atol: float = 1e-5,
+) -> torch.Tensor:
+    """Run CPU ONNX Runtime and compare its raw logits with PyTorch."""
+    try:
+        import numpy as np
+        import onnxruntime as ort
+    except ImportError as error:
+        raise ImportError(
+            "ONNX Runtime validation requires the 'onnx' extra. Install with: "
+            "uv sync --extra root-gnn --extra onnx"
+        ) from error
+    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
+    values = {
+        name: tensor.detach().cpu().numpy()
+        for name, tensor in zip(_INPUT_NAMES, inputs.as_tuple())
+    }
+    result = torch.from_numpy(np.asarray(session.run(["logits"], values)[0]))
+    torch.testing.assert_close(result, reference.detach().cpu(), rtol=rtol, atol=atol)
+    return result
+
+
+def _metadata(
+    model: nn.Module, *, checkpoint: str | None, task_config: Mapping[str, Any] | None
+) -> dict[str, Any]:
+    config = model.checkpoint_config()
+    return {
+        "model_family": "root_gnn",
+        "model_config": config,
+        "task_config": dict(task_config or {}),
+        "feature_schema_version": FEATURE_SCHEMA_VERSION,
+        "graph_schema_version": GRAPH_SCHEMA_VERSION,
+        "onnx_opset": ONNX_OPSET,
+        "input_names": list(_INPUT_NAMES),
+        "output_name": "logits",
+        "checkpoint": (Path(checkpoint).name if checkpoint else None),
+        "checkpoint_sha256": (
+            hashlib.sha256(Path(checkpoint).read_bytes()).hexdigest()
+            if checkpoint
+            else None
+        ),
+    }
+
+
+def export_root_gnn_onnx(
+    model: nn.Module,
+    example_batch: Any,
+    output_path: str | os.PathLike[str],
+    *,
+    opset: int = ONNX_OPSET,
+    task_config: Mapping[str, Any] | None = None,
+    checkpoint: str | None = None,
+    overwrite: bool = False,
+) -> Path:
+    """Export and validate a ROOT-GNN model from an already prepared batch."""
+    if opset != ONNX_OPSET:
+        raise ValueError(f"only ONNX opset {ONNX_OPSET} is supported")
+    _optional_onnx()
+    inputs = (
+        example_batch
+        if isinstance(example_batch, RootGNNExportInputs)
+        else inputs_from_graph_batch(example_batch)
+    )
+    adapter = RootGNNExportAdapter(model).eval()
+    output = Path(output_path)
+    if output.suffix.lower() != ".onnx":
+        raise ValueError("export output must have a .onnx extension")
+    if output.exists() and not overwrite:
+        raise FileExistsError(f"export output already exists: {output}")
+    output.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.TemporaryDirectory(
+        prefix=f".{output.stem}-", dir=output.parent
+    ) as temp_dir:
+        temporary = Path(temp_dir) / output.name
+        torch.onnx.export(
+            adapter,
+            inputs.as_tuple(),
+            str(temporary),
+            input_names=list(_INPUT_NAMES),
+            output_names=["logits"],
+            opset_version=opset,
+            dynamic_axes={
+                "node_features": {0: "nodes"},
+                "edge_features": {0: "edges"},
+                "edge_src": {0: "edges"},
+                "edge_dst": {0: "edges"},
+                "node_batch": {0: "nodes"},
+                "global_features": {0: "graphs"},
+                "logits": {0: "graphs"},
+            },
+            do_constant_folding=True,
+        )
+        validate_onnx(temporary)
+        with torch.inference_mode():
+            reference = adapter(*inputs.as_tuple())
+        validate_onnx_runtime(temporary, inputs, reference)
+        temporary.replace(output)
+    metadata_path = output.with_suffix(output.suffix + ".json")
+    metadata_path.write_text(
+        json.dumps(
+            _metadata(model, checkpoint=checkpoint, task_config=task_config),
+            indent=2,
+            sort_keys=True,
+        )
+        + "\n"
+    )
+    return output
+
+
+def _model_from_checkpoint(payload: Mapping[str, Any], example_batch: Any) -> nn.Module:
+    config = dict(payload.get("model_config") or {})
+    if str(config.get("family", "root_gnn")) != "root_gnn":
+        raise ValueError("checkpoint model family is not root_gnn")
+    graph = example_batch.graph if hasattr(example_batch, "graph") else None
+    if graph is None:
+        raise TypeError("checkpoint reconstruction needs a GraphBatch example")
+    kwargs = {
+        key: config[key]
+        for key in ("hid_size", "n_layers", "n_proc_steps", "dropout")
+        if key in config
+    }
+    model_class = str(config.get("class", "EdgeNetwork"))
+    if "FineTuned" in model_class or str(config.get("name", "")).startswith("fine"):
+        backbone = EdgeNetwork(
+            graph, getattr(example_batch, "global_features", None), out_size=1, **kwargs
+        )
+        model = FineTunedEdgeNetwork(
+            backbone,
+            int(config.get("out_size", 1)),
+            freeze_backbone=bool(config.get("freeze_backbone", False)),
+        )
+    else:
+        model = EdgeNetwork(
+            graph,
+            getattr(example_batch, "global_features", None),
+            out_size=int(config.get("out_size", 1)),
+            **kwargs,
+        )
+    load_model_weights(model, payload)
+    return model
+
+
+def export_checkpoint_to_onnx(
+    checkpoint_path: str | os.PathLike[str],
+    output_path: str | os.PathLike[str],
+    *,
+    example_batch: Any,
+    opset: int = ONNX_OPSET,
+    overwrite: bool = False,
+) -> Path:
+    """Load a Task 11 checkpoint and export its reconstructed ROOT-GNN model."""
+    payload = CheckpointManager.load(checkpoint_path, map_location="cpu")
+    model = _model_from_checkpoint(payload, example_batch)
+    return export_root_gnn_onnx(
+        model,
+        example_batch,
+        output_path,
+        opset=opset,
+        task_config=payload.get("task_config"),
+        checkpoint=str(checkpoint_path),
+        overwrite=overwrite,
+    )

src/gnn4colliders/inference/__init__.py

@@ -0,0 +1,15 @@
+"""Prediction, evaluation, and output writing infrastructure."""
+
+from .predictor import Predictor, load_model_and_task_for_inference
+from .results import EvaluationResult, PredictionResult
+from .root_writer import write_root_scores
+from .writers import write_npz
+
+__all__ = [
+    "EvaluationResult",
+    "PredictionResult",
+    "Predictor",
+    "load_model_and_task_for_inference",
+    "write_npz",
+    "write_root_scores",
+]

src/gnn4colliders/inference/predictor.py

@@ -0,0 +1,231 @@
+"""Model execution for prediction and full-split evaluation."""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable, Mapping
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+from torch import nn
+
+from gnn4colliders.data.graph_dataset import GraphBatch
+from gnn4colliders.distributed import (
+    DistributedContext,
+    gather_objects,
+    gather_tensor,
+    prepare_model,
+)
+from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask
+from gnn4colliders.training.checkpoint import CheckpointManager, load_model_weights
+
+from .results import EvaluationResult, PredictionResult
+
+
+def _forward(model: nn.Module, batch: Any) -> torch.Tensor:
+    """Call the public GraphBatch or graph/global model boundary."""
+    if isinstance(batch, GraphBatch):
+        return model(batch.graph, batch.global_features)
+    try:
+        signature = inspect.signature(model.forward)
+        positional = [
+            parameter
+            for parameter in signature.parameters.values()
+            if parameter.kind
+            in (parameter.POSITIONAL_ONLY, parameter.POSITIONAL_OR_KEYWORD)
+        ]
+    except (TypeError, ValueError):
+        positional = []
+    if len(positional) <= 1:
+        return model(batch)
+    if hasattr(batch, "graph"):
+        return model(batch.graph, getattr(batch, "global_features", None))
+    raise TypeError("model requires graph/global_features, but batch has no graph")
+
+
+def _batch_metadata(batch: Any) -> Any:
+    return getattr(batch, "metadata", None)
+
+
+def _labels(batch: Any) -> torch.Tensor | None:
+    labels = getattr(batch, "labels", None)
+    return None if labels is None else torch.as_tensor(labels).detach().cpu()
+
+
+class Predictor:
+    """Run a task-aware model over an ordered loader."""
+
+    def __init__(
+        self,
+        model: nn.Module,
+        task: Any,
+        *,
+        device: torch.device | str = "cpu",
+        distributed_context: DistributedContext | None = None,
+    ) -> None:
+        self.distributed_context = (
+            distributed_context or DistributedContext.single_process(device)
+        )
+        self.device = self.distributed_context.device
+        self.model = prepare_model(model, self.distributed_context)
+        self.task = task
+
+    def predict(self, loader: Any) -> PredictionResult:
+        self.model.eval()
+        logits_parts: list[torch.Tensor] = []
+        label_parts: list[torch.Tensor] = []
+        sample_ids: list[str] = []
+        metadata_parts: list[Any] = []
+        has_labels = True
+
+        with torch.inference_mode():
+            for batch in loader:
+                moved = batch.to(self.device) if hasattr(batch, "to") else batch
+                logits = _forward(self.model, moved)
+                if not isinstance(logits, torch.Tensor):
+                    raise TypeError("model must return a torch.Tensor of logits")
+                logits_parts.append(logits.detach().cpu())
+                labels = _labels(moved)
+                if labels is None:
+                    has_labels = False
+                else:
+                    label_parts.append(labels)
+                metadata = _batch_metadata(moved)
+                if metadata is not None:
+                    ids = getattr(metadata, "sample_id", None)
+                    if ids is not None:
+                        sample_ids.extend(str(value) for value in ids)
+                    metadata_parts.append(metadata)
+
+        if not logits_parts:
+            raise ValueError("cannot predict on an empty loader")
+        logits = torch.cat(logits_parts, dim=0)
+        if self.distributed_context.enabled:
+            logits = gather_tensor(logits, self.distributed_context)
+            if label_parts:
+                labels_local = torch.cat(label_parts, dim=0)
+                labels_result = gather_tensor(labels_local, self.distributed_context)
+            else:
+                labels_result = None
+            ids = tuple(
+                value
+                for group in gather_objects(tuple(sample_ids), self.distributed_context)
+                for value in group
+            )
+            local_metadata = _combine_metadata(metadata_parts)
+            if local_metadata is not None and hasattr(local_metadata, "weight"):
+                metadata = SimpleNamespace(
+                    fold=gather_tensor(local_metadata.fold, self.distributed_context),
+                    weight=gather_tensor(
+                        local_metadata.weight, self.distributed_context
+                    ),
+                    sample_id=ids,
+                )
+            sample_ids = list(ids)
+        else:
+            labels_result = torch.cat(label_parts, dim=0) if has_labels else None
+        task_output = self.task.predictions(logits)
+        if not self.distributed_context.enabled:
+            metadata = _combine_metadata(metadata_parts)
+        if not sample_ids and metadata is not None:
+            sample_ids = [str(value) for value in getattr(metadata, "sample_id", ())]
+        if len(sample_ids) != logits.shape[0]:
+            raise ValueError(
+                "loader metadata must provide one sample_id for every prediction"
+            )
+        return PredictionResult(
+            sample_ids=tuple(sample_ids),
+            logits=logits,
+            scores=task_output["scores"].detach().cpu(),
+            predictions=task_output["predictions"].detach().cpu(),
+            labels=labels_result,
+            metadata=metadata,
+            extra=getattr(metadata, "extra", {}) if metadata is not None else {},
+        )
+
+    def evaluate(self, loader: Any) -> EvaluationResult:
+        result = self.predict(loader)
+        if result.labels is None:
+            raise ValueError("evaluation requires labels")
+        weights = result.weights
+        batch = SimpleNamespace(
+            labels=result.labels,
+            metadata=SimpleNamespace(weight=weights) if weights is not None else None,
+        )
+        if weights is None:
+            raise ValueError("evaluation requires named metadata.weight")
+        metrics = {
+            key: float(value)
+            for key, value in self.task.metrics(result.logits, batch).items()
+        }
+        return EvaluationResult(predictions=result, metrics=metrics)
+
+
+def _combine_metadata(parts: list[Any]) -> Any | None:
+    if not parts:
+        return None
+    first = parts[0]
+    cls = type(first)
+    fields = ("fold", "weight", "sample_id", "extra")
+    values: dict[str, Any] = {}
+    for name in fields:
+        if not hasattr(first, name):
+            continue
+        current = [getattr(part, name) for part in parts]
+        if name in {"fold", "weight"}:
+            values[name] = torch.cat(
+                [torch.as_tensor(value).cpu() for value in current]
+            )
+        elif name == "sample_id":
+            values[name] = tuple(value for group in current for value in group)
+        elif name == "extra":
+            keys = sorted({key for group in current for key in group})
+            values[name] = {
+                key: tuple(value for group in current for value in group.get(key, ()))
+                for key in keys
+            }
+    try:
+        return cls(**values)
+    except TypeError:
+        return SimpleNamespace(**values)
+
+
+def load_model_and_task_for_inference(
+    checkpoint: str,
+    *,
+    model_factory: Callable[[Mapping[str, Any]], nn.Module],
+    task_factory: Callable[[Mapping[str, Any]], Any] | None = None,
+    map_location: torch.device | str = "cpu",
+) -> tuple[nn.Module, Any]:
+    """Build a model from checkpoint metadata and load weights only.
+
+    Factories remain caller-owned because ROOT-GNN construction needs a sample
+    graph.  The helper deliberately never restores optimizer or lifecycle state.
+    """
+    checkpoint_data = CheckpointManager.load(checkpoint, map_location=map_location)
+    model_config = checkpoint_data.get("model_config")
+    if not isinstance(model_config, Mapping):
+        raise ValueError("checkpoint does not contain reconstructable model_config")
+    model = model_factory(model_config)
+    load_model_weights(model, checkpoint_data)
+    task_config = checkpoint_data.get("task_config") or {}
+    if task_factory is not None:
+        task = task_factory(task_config)
+    else:
+        task_type = str(task_config.get("type", "")).lower()
+        if "multi" in task_type:
+            task = MulticlassClassificationTask(
+                **_task_kwargs(task_config, "absolute_weights")
+            )
+        elif "binary" in task_type:
+            task = BinaryClassificationTask(
+                **_task_kwargs(task_config, "absolute_weights", "threshold")
+            )
+        else:
+            raise ValueError("checkpoint does not contain a supported task_config")
+    return model, task
+
+
+def _task_kwargs(config: Mapping[str, Any], *names: str) -> dict[str, Any]:
+    return {name: config[name] for name in names if name in config}

src/gnn4colliders/inference/results.py

@@ -0,0 +1,52 @@
+"""Typed outputs produced by model inference."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+
+@dataclass(frozen=True)
+class PredictionResult:
+    """Predictions in the exact order in which the loader yielded events.
+
+    Tensor fields are CPU tensors.  Keeping the tensors detached and on the
+    CPU makes a result safe to retain after prediction and avoids tying its
+    lifetime to a CUDA context.
+    """
+
+    sample_ids: tuple[str, ...]
+    logits: torch.Tensor
+    scores: torch.Tensor
+    predictions: torch.Tensor
+    labels: torch.Tensor | None = None
+    metadata: Any | None = None
+    extra: Mapping[str, Sequence[Any]] = field(default_factory=dict)
+
+    @property
+    def fold(self) -> torch.Tensor | None:
+        return _metadata_value(self.metadata, "fold")
+
+    @property
+    def weights(self) -> torch.Tensor | None:
+        return _metadata_value(self.metadata, "weight")
+
+
+@dataclass(frozen=True)
+class EvaluationResult:
+    """A prediction result and metrics calculated over the complete split."""
+
+    predictions: PredictionResult
+    metrics: Mapping[str, float]
+
+
+def _metadata_value(metadata: Any, name: str) -> torch.Tensor | None:
+    if metadata is None:
+        return None
+    value = getattr(metadata, name, None)
+    if value is None and isinstance(metadata, Mapping):
+        value = metadata.get(name)
+    return None if value is None else torch.as_tensor(value).cpu()

src/gnn4colliders/inference/root_writer.py

@@ -0,0 +1,79 @@
+"""Optional ROOT score output, kept separate from model execution."""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+import numpy as np
+
+from .results import PredictionResult
+
+_ENTRY_SUFFIX = re.compile(r":(\d+)$")
+
+
+def write_root_scores(
+    result: PredictionResult,
+    *,
+    input_path: str | Path,
+    output_path: str | Path,
+    tree_name: str,
+    score_prefix: str = "score",
+) -> Path:
+    """Clone a ROOT tree's columns and add aligned score branches.
+
+    Sample IDs ending in ``:<entry>`` are aligned by source entry.  If IDs do
+    not carry an entry suffix, the result must contain one score per source
+    entry; this makes accidental misalignment fail loudly.  Unselected source
+    entries receive NaN scores and ``selection_pass=0``.
+    """
+    try:
+        import uproot
+    except ImportError as error:  # pragma: no cover - optional dependency
+        raise ImportError("write_root_scores requires uproot") from error
+
+    source = Path(input_path)
+    target = Path(output_path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    with uproot.open(source) as root_file:
+        tree = root_file[tree_name]
+        columns = tree.arrays(library="np")
+        entries = int(tree.num_entries)
+
+    indices: list[int] = []
+    for sample_id in result.sample_ids:
+        match = _ENTRY_SUFFIX.search(sample_id)
+        if match is None:
+            if len(result.sample_ids) != entries:
+                raise ValueError(
+                    "ROOT alignment requires sample IDs ending in ':<entry>' "
+                    "when predictions do not cover every source entry"
+                )
+            indices = list(range(entries))
+            break
+        indices.append(int(match.group(1)))
+    if len(indices) != len(result.sample_ids) or any(
+        index < 0 or index >= entries for index in indices
+    ):
+        raise ValueError("prediction sample IDs contain invalid ROOT entry indices")
+
+    selection_pass = np.zeros(entries, dtype=np.int32)
+    selection_pass[indices] = 1
+    output = dict(columns)
+    output["selection_pass"] = selection_pass
+    scores = np.asarray(result.scores)
+    if scores.ndim == 1:
+        branch = np.full(entries, np.nan, dtype=np.float32)
+        branch[indices] = scores.astype(np.float32, copy=False)
+        output[score_prefix] = branch
+    elif scores.ndim == 2:
+        for class_index in range(scores.shape[1]):
+            branch = np.full(entries, np.nan, dtype=np.float32)
+            branch[indices] = scores[:, class_index].astype(np.float32, copy=False)
+            output[f"{score_prefix}_class_{class_index}"] = branch
+    else:
+        raise ValueError("scores must have shape [N] or [N, C]")
+
+    with uproot.recreate(target) as root_file:
+        root_file[tree_name] = output
+    return target

src/gnn4colliders/inference/writers.py

@@ -0,0 +1,32 @@
+"""Serialization of prediction results."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from .results import PredictionResult
+
+
+def write_npz(result: PredictionResult, path: str | Path) -> Path:
+    """Write named prediction fields to a compressed NumPy archive."""
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    values: dict[str, Any] = {
+        "sample_id": np.asarray(result.sample_ids, dtype=str),
+        "logits": result.logits.numpy(),
+        "scores": result.scores.numpy(),
+        "predictions": result.predictions.numpy(),
+    }
+    if result.labels is not None:
+        values["labels"] = result.labels.numpy()
+    for name, attribute in (("fold", "fold"), ("weight", "weights")):
+        value = getattr(result, attribute)
+        if value is not None:
+            values[name] = value.numpy()
+    for name, value in result.extra.items():
+        values[name] = np.asarray(value)
+    np.savez_compressed(target, **values)
+    return target

tests/unit/export/test_onnx_export.py

@@ -0,0 +1,62 @@
+import pytest
+import torch
+
+pytest.importorskip("dgl")
+
+from gnn4colliders.data import EventMetadata, GraphSample, batch_graph_samples
+from gnn4colliders.export import (
+    RootGNNExportAdapter,
+    export_root_gnn_onnx,
+    inputs_from_graph_batch,
+)
+from gnn4colliders.graphs import build_dgl_graph
+from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
+
+
+def _batch():
+    samples = []
+    for index, count in enumerate((2, 3)):
+        nodes = torch.arange(count * 3, dtype=torch.float32).reshape(count, 3)
+        samples.append(
+            GraphSample(
+                build_dgl_graph(nodes),
+                torch.tensor(index),
+                torch.tensor([1.0, 2.0]),
+                EventMetadata(index, 1.0, f"fixture:{index}"),
+            )
+        )
+    return batch_graph_samples(samples)
+
+
+def test_export_inputs_preserve_graph_membership_and_edges():
+    batch = _batch()
+    inputs = inputs_from_graph_batch(batch)
+    src, _ = batch.graph.edges(order="eid")
+    assert torch.equal(inputs.node_features, batch.graph.ndata["features"])
+    assert torch.equal(inputs.edge_features, batch.graph.edata["features"])
+    assert torch.equal(inputs.edge_src, src)
+    assert inputs.node_batch.tolist() == [0, 0, 1, 1, 1]
+    assert inputs.global_features.shape == (2, 2)
+
+
+@pytest.mark.parametrize("fine_tuned", [False, True])
+def test_tensor_adapter_matches_native_model(fine_tuned):
+    batch = _batch()
+    model = EdgeNetwork(batch.graph, batch.global_features, 8, 3, 2, 2).eval()
+    if fine_tuned:
+        model = FineTunedEdgeNetwork.from_pretrained(model, 1).eval()
+    inputs = inputs_from_graph_batch(batch)
+    with torch.inference_mode():
+        native = model(batch.graph, batch.global_features)
+        adapted = RootGNNExportAdapter(model)(*inputs.as_tuple())
+    torch.testing.assert_close(adapted, native, rtol=1e-5, atol=1e-5)
+
+
+def test_onnx_export_is_optional_and_writes_validated_model(tmp_path):
+    pytest.importorskip("onnx")
+    pytest.importorskip("onnxruntime")
+    batch = _batch()
+    model = EdgeNetwork(batch.graph, batch.global_features, 8, 3, 2, 2).eval()
+    output = export_root_gnn_onnx(model, batch, tmp_path / "model.onnx")
+    assert output.exists()
+    assert output.with_suffix(".onnx.json").exists()

tests/unit/inference/test_predictor.py

@@ -0,0 +1,83 @@
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+from torch import nn
+
+from gnn4colliders.data import BatchMetadata
+from gnn4colliders.inference import Predictor, write_npz
+from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask
+
+
+class _Model(nn.Module):
+    def __init__(self, outputs):
+        super().__init__()
+        self.outputs = torch.as_tensor(outputs, dtype=torch.float32)
+        self.dropout = nn.Dropout(0.9)
+
+    def forward(self, batch):
+        count = (
+            batch.labels.shape[0]
+            if hasattr(batch, "labels")
+            else len(batch.metadata.sample_id)
+        )
+        return self.outputs[batch.offset : batch.offset + count]
+
+
+def _loader(outputs, labels, *, offset=0, ids=("a", "b")):
+    metadata = BatchMetadata(
+        fold=torch.tensor([0, 1]),
+        weight=torch.tensor([1.0, 2.0]),
+        sample_id=ids,
+        extra={"source_file": ("events.root", "events.root")},
+    )
+    return [
+        SimpleNamespace(
+            labels=torch.tensor(labels),
+            metadata=metadata,
+            offset=offset,
+            to=lambda device: SimpleNamespace(
+                labels=torch.tensor(labels), metadata=metadata, offset=offset
+            ),
+        )
+    ]
+
+
+def test_predictor_preserves_named_metadata_and_uses_task_scores():
+    loader = _loader([[-2.0], [2.0]], [0, 1])
+    result = Predictor(_Model([[-2.0], [2.0]]), BinaryClassificationTask()).predict(
+        loader
+    )
+    assert result.sample_ids == ("a", "b")
+    assert result.scores.tolist() == [
+        torch.sigmoid(torch.tensor(-2.0)).item(),
+        torch.sigmoid(torch.tensor(2.0)).item(),
+    ]
+    assert result.predictions.tolist() == [False, True]
+    assert result.weights.tolist() == [1.0, 2.0]
+    assert result.extra["source_file"] == ("events.root", "events.root")
+
+
+def test_predictor_supports_unlabeled_multiclass_batches():
+    metadata = BatchMetadata(
+        fold=torch.tensor([0]), weight=torch.tensor([1.0]), sample_id=("x",)
+    )
+    batch = SimpleNamespace(metadata=metadata, offset=0, to=lambda device: batch)
+    result = Predictor(
+        _Model([[1.0, 3.0, 2.0]]), MulticlassClassificationTask()
+    ).predict([batch])
+    assert result.labels is None
+    assert result.scores.shape == (1, 3)
+    assert result.predictions.tolist() == [1]
+
+
+def test_npz_round_trip(tmp_path):
+    result = Predictor(_Model([[-1.0], [1.0]]), BinaryClassificationTask()).predict(
+        _loader([[-1.0], [1.0]], [0, 1])
+    )
+    path = write_npz(result, tmp_path / "predictions.npz")
+    with np.load(path) as values:
+        assert values["sample_id"].tolist() == ["a", "b"]
+        np.testing.assert_array_equal(values["logits"], result.logits.numpy())
+        np.testing.assert_array_equal(values["labels"], result.labels.numpy())
+        np.testing.assert_array_equal(values["fold"], result.fold.numpy())