fix: enable validated ONNX export

pyproject.toml

@@ -36,6 +36,8 @@ dev = [
     "pytest-cov>=5",
     "ruff>=0.6",
     "twine>=6,<7",
+    "onnx>=1.16,<2",
+    "onnxruntime>=1.18,<2",
 ]
 
 [tool.setuptools.packages.find]

src/gnn4colliders/export/onnx.py

@@ -106,23 +106,30 @@ def inputs_from_graph_batch(batch: Any) -> RootGNNExportInputs:
 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
+    membership = _group_membership(groups, count, values)
+    result = membership @ values
+    sizes = membership.sum(dim=1, keepdim=True)
+    return result / sizes.clamp_min(1)
 
 
 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
+    return _group_membership(groups, count, values) @ values
+
+
+def _group_membership(
+    groups: torch.Tensor, count: int, values: torch.Tensor
+) -> torch.Tensor:
+    """Return a dense group-membership matrix for export-safe reduction.
+
+    ``index_add`` is efficient in native PyTorch but exports to ONNX scatter
+    operations whose behavior for duplicate indices is not reliable.  A
+    boolean membership matrix has unambiguous reduction semantics for the
+    repeated node/edge group IDs used by batched graphs.
+    """
+    labels = torch.arange(count, device=groups.device).unsqueeze(1)
+    return (labels == groups.unsqueeze(0)).to(dtype=values.dtype)
 
 
 class RootGNNExportAdapter(nn.Module):
@@ -169,9 +176,7 @@ class RootGNNExportAdapter(nn.Module):
                     1,
                 )
             )
-            node_messages = node_h.new_zeros(node_h.shape).index_add(
-                0, edge_dst, edge_h
-            )
+            node_messages = _sum_by_group(edge_h, edge_dst, node_h.shape[0])
             node_h = self.backbone.node_update(
                 torch.cat((node_h, node_messages, global_h[node_batch]), 1)
             )

uv.lock

@@ -239,6 +239,8 @@ root-gnn = [
 [package.dev-dependencies]
 dev = [
     { name = "matplotlib", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "onnx", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+    { name = "onnxruntime", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "pytest", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "pytest-cov", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "ruff", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
@@ -262,6 +264,8 @@ provides-extras = ["root-gnn", "onnx"]
 [package.metadata.requires-dev]
 dev = [
     { name = "matplotlib", specifier = ">=3.8,<4" },
+    { name = "onnx", specifier = ">=1.16,<2" },
+    { name = "onnxruntime", specifier = ">=1.18,<2" },
     { name = "pytest", specifier = ">=8" },
     { name = "pytest-cov", specifier = ">=5" },
     { name = "ruff", specifier = ">=0.6" },