ci: support macOS core development

.github/workflows/ci.yml

@@ -34,6 +34,18 @@ jobs:
       - run: uv sync --dev
       - run: uv run pytest
 
+  test-macos:
+    name: test (macOS ROOT-GNN CPU)
+    runs-on: macos-14
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.8.x"
+          enable-cache: true
+      - run: uv sync --dev --extra root-gnn
+      - run: uv run pytest tests/unit tests/integration
+
   package:
     name: package
     runs-on: ubuntu-latest

README.md

@@ -19,20 +19,50 @@ historical checkpoint investigation, not a supported runtime backend.
 
 ## Installation
 
-The supported development environment is Python 3.12 (`>=3.12,<3.13`), with
-PyTorch 2.2.2 and the optional ROOT-GNN stack DGL 2.4.0. The canonical setup
-is:
+The supported development environment is Python 3.12 (`>=3.12,<3.13`). Core
+development is supported on macOS and Linux:
 
 ```bash
+# macOS (Apple Silicon): CPU ROOT-GNN development and tests
+uv sync --dev --extra root-gnn
+
+# Linux x86_64 with an NVIDIA GPU: validated ROOT-GNN development
 uv sync --dev --extra root-gnn
 ```
 
 The core package can be installed without DGL when only shared data or task
 code is needed. ROOT-GNN models, graph construction, and ROOT-GNN parity tests
-require the `root-gnn` extra. The DGL extra uses the validated CUDA 12.1 wheel
-source configured in `pyproject.toml`; a compatible NVIDIA driver is still
-required. Do not add site-specific CUDA, Slurm, or filesystem paths to model
-or task configuration.
+require the `root-gnn` extra. On Linux x86_64, it uses the validated CUDA 12.1
+wheels configured in `pyproject.toml`; a compatible NVIDIA driver is still
+required. On Apple Silicon macOS, it installs the CPU DGL wheel, supporting
+local graph/cache development. The default ROOT-GNN backend performs training
+with native PyTorch graph tensors, so it runs on Apple MPS, NVIDIA CUDA, and
+CPU; DGL remains a cache and legacy-compatibility adapter. Do not add
+site-specific CUDA, Slurm, or filesystem paths to model or task configuration.
+
+Use the MPS profile on an Apple Silicon Mac:
+
+```bash
+uv run gnn4colliders train environment=macos
+```
+
+## Data samples
+
+ROOT inputs are available from the
+[HWresearch/Delphes dataset](https://huggingface.co/datasets/HWresearch/Delphes).
+Download the 64-event smoke-test sample with the Hugging Face CLI:
+
+```bash
+hf download HWresearch/Delphes testing/ttH_NLO_64.root \
+  --repo-type dataset --local-dir data/raw
+```
+
+The sample is `data/raw/testing/ttH_NLO_64.root`, has tree name `output`, and
+is suitable for checking the prepare/train workflow. The dataset also provides
+larger process-specific ROOT samples under `samples/`, derived datasets under
+`derived/`, and analysis-specific ntuples under `analyses/`. These data are
+intentionally ignored by Git; inspect a selected ROOT file's tree and branches
+before writing its preparation configuration.
 
 ## Quick start
 

docs/architecture.md

@@ -28,7 +28,7 @@ shared collider feature construction
 | `data` | ROOT/Awkward ingestion, event samples, metadata, graph caches, folds, and batching |
 | `features` | Shared collider-object features and derived physics quantities |
 | `graphs` | Topology, edge features, and the DGL representation adapter |
-| `models/root_gnn` | ROOT-GNN encoders, message passing, classifier, and transfer boundary |
+| `models/root_gnn` | ROOT-GNN encoders, portable tensor message passing, classifier, and transfer boundary |
 | `tasks` | Loss, score/prediction, labels, weights, and metrics |
 | `training` | Optimizer lifecycle, validation, early stopping, checkpointing, and reproducibility |
 | `inference` | Ordered prediction/evaluation and NPZ/ROOT output adapters |
@@ -36,8 +36,12 @@ shared collider feature construction
 | `config` / `cli` | Semantic Hydra composition and thin user-facing commands |
 
 `EventSample` is shared infrastructure, not a ROOT-GNN object. `GraphSample`
-is the current representation-specific adapter. This separation is the
-extension point for a future sequence/token representation.
+is the current representation-specific adapter. At batching, the default
+ROOT-GNN backend converts cached DGL graphs into `TensorGraph` values with
+explicit node/edge indices and graph membership. Native message passing then
+uses only PyTorch tensor operations, allowing the same model to train on CPU,
+CUDA, and Apple MPS. This separation is the extension point for a future
+sequence/token representation.
 
 The new public metadata contract is named `EventMetadata(fold, weight,
 sample_id, extra)`. The legacy positional tracking tensor is accepted only by
@@ -46,9 +50,8 @@ Level-2 graph cache; replacing it with a universal cache would couple future
 model families to DGL.
 
 For deployment, a prepared `GraphBatch` can pass through the isolated
-`RootGNNExportAdapter` into an ONNX model. This is an inference boundary only;
-the native model continues to consume DGL graphs and ONNX does not read ROOT
-or construct collider features.
+`RootGNNExportAdapter` into an ONNX model. ONNX is still an inference boundary:
+it does not read ROOT or construct collider features.
 
 ## Current public workflow
 

docs/configuration.md

@@ -88,7 +88,9 @@ to load weights into a new task head; these options are mutually exclusive.
 
 ## Environment profiles
 
-`environment=local` selects CPU by default. `environment=perlmutter` selects
+`environment=local` selects CPU by default. `environment=macos` selects the
+Apple MPS device for the portable PyTorch ROOT-GNN backend.
+`environment=perlmutter` selects
 the CUDA device and a conventional output-root pattern. Profiles should hold
 device/output policy only; site-specific module loads and filesystem paths
 belong in a launcher or shell environment.

docs/migration.md

@@ -104,6 +104,13 @@ active bug when nonempty globals are supplied (`Pretrained_Output` ignores its
 argument); parity therefore characterizes its supported no-global path, while
 the rewritten model supports both global and fallback modes.
 
+Task 22 adds `TensorEdgeNetwork`, a native PyTorch realization of the same
+message-passing equations. It receives explicit node/edge index tensors from
+the batching boundary, retains parameter names for legacy checkpoint loading,
+and has fixed-weight CPU parity coverage against the DGL backend. It is the
+default training backend for CPU, CUDA, and Apple MPS; DGL remains available as
+a cache and compatibility adapter during the transition.
+
 Port `models.GCN.Edge_Network` first. Preserve constructor parameters,
 `forward(graph, global_feats)`, feature keys, processor order, MLP LayerNorm
 placement, and logits shape. Compare intermediate and final tensors on fixed

docs/testing.md

@@ -19,6 +19,10 @@ uv run pytest -m gpu -v
 GNN4COLLIDERS_ROOT_FIXTURE=/path/to/reduced.root uv run pytest -m real_data -v
 ```
 
+On Apple Silicon macOS, `uv sync --dev --extra root-gnn` supports CPU
+ROOT-GNN unit and integration testing. CUDA and the required ROOT-GNN parity
+gate are validated on Linux x86_64 with the `root-gnn` extra installed.
+
 Unit tests use deterministic, small tensors and generated ROOT files. DGL,
 ONNX, CUDA, distributed execution, and reduced real-data fixtures remain
 optional layers. Tests should assert public contracts and scientific

pyproject.toml

@@ -19,7 +19,10 @@ dependencies = [
 
 [project.optional-dependencies]
 root-gnn = [
-    "dgl==2.4.0+cu121",
+    "dgl==2.4.0+cu121; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "dgl==2.2.0; sys_platform == 'darwin' and platform_machine == 'arm64'",
+    "pydantic>=2,<3; sys_platform == 'darwin' and platform_machine == 'arm64'",
+    "torchdata>=0.7,<0.8; sys_platform == 'darwin' and platform_machine == 'arm64'",
 ]
 onnx = [
     "onnx>=1.16,<2",
@@ -50,12 +53,11 @@ version = { attr = "gnn4colliders.__version__" }
 "gnn4colliders.configs" = ["*.yaml", "*/*.yaml", "*/*/*.yaml"]
 
 [tool.uv.sources]
-torch = { index = "pytorch-cu121" }
-dgl = { index = "dgl" }
-
-[tool.uv]
-environments = [
-    "sys_platform == 'linux' and platform_machine == 'x86_64' and python_full_version >= '3.12' and python_full_version < '3.13'",
+torch = [
+    { index = "pytorch-cu121", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
+]
+dgl = [
+    { index = "dgl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
 ]
 
 [[tool.uv.index]]

src/gnn4colliders/config/application.py

@@ -72,6 +72,7 @@ def loaders(
         validation_folds=frozenset(config.data.splits.validation_folds),
         test_folds=frozenset(config.data.splits.test_folds),
     )
+    tensor_graphs = str(config.model.get("backend", "torch")).lower() == "torch"
     result = {}
     for name in ("train", "validation", "test"):
         selected = GraphDataset(select_split(samples.samples, split, name))
@@ -87,6 +88,7 @@ def loaders(
                 context,
                 shuffle=bool(config.data.shuffle) if name == "train" else False,
                 seed=int(config.data.seed),
+                tensor_graphs=tensor_graphs,
             )
             if loader_type is DistributedGraphDataLoader
             else loader_type(
@@ -94,6 +96,7 @@ def loaders(
                 int(config.data.batch_size),
                 shuffle=bool(config.data.shuffle) if name == "train" else False,
                 seed=int(config.data.seed),
+                tensor_graphs=tensor_graphs,
             )
         )
     return result

src/gnn4colliders/config/factories.py

@@ -8,7 +8,11 @@ from typing import Any
 import torch
 
 from gnn4colliders.distributed import DistributedContext
-from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
+from gnn4colliders.models.root_gnn import (
+    EdgeNetwork,
+    FineTunedEdgeNetwork,
+    TensorEdgeNetwork,
+)
 from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask
 from gnn4colliders.training import (
     EarlyStopping,
@@ -40,6 +44,10 @@ def build_model(
         )
     if family != "root_gnn":
         raise ValueError(f"unsupported model family {family!r}")
+    backend = str(_get(config, "backend", "torch")).lower()
+    if backend not in {"torch", "dgl"}:
+        raise ValueError(f"unsupported root_gnn backend {backend!r}")
+    model_class = TensorEdgeNetwork if backend == "torch" else EdgeNetwork
     kwargs = {
         "hid_size": int(_get(config, "hid_size", _get(config, "hidden_dim", 128))),
         "out_size": int(_get(config, "out_size", 1)),
@@ -50,7 +58,7 @@ def build_model(
         "dropout": float(_get(config, "dropout", 0.0)),
     }
     if name in {"edge_network", "edge"}:
-        return EdgeNetwork(sample_graph, sample_global, **kwargs)
+        return model_class(sample_graph, sample_global, **kwargs)
     if name in {"fine_tuned_edge_network", "finetuned_edge_network", "fine_tuned"}:
         checkpoint = _get(config, "pretrained_checkpoint")
         if not checkpoint:
@@ -65,7 +73,7 @@ def build_model(
         base_config.update(
             {key: value for key, value in kwargs.items() if key != "out_size"}
         )
-        backbone = EdgeNetwork(
+        backbone = model_class(
             sample_graph,
             sample_global,
             hid_size=int(base_config.get("hid_size", kwargs["hid_size"])),

src/gnn4colliders/configs/environment/macos.yaml

@@ -0,0 +1,2 @@
+device: mps
+output_root: outputs

src/gnn4colliders/configs/model/root_gnn/edge_network.yaml

@@ -1,4 +1,5 @@
 family: root_gnn
+backend: torch
 name: edge_network
 hid_size: 128
 out_size: 12

src/gnn4colliders/configs/model/root_gnn/fine_tuned_edge_network.yaml

@@ -1,4 +1,5 @@
 family: root_gnn
+backend: torch
 name: fine_tuned_edge_network
 hid_size: 128
 out_size: 1

src/gnn4colliders/data/__init__.py

@@ -8,6 +8,7 @@ from .graph_dataset import (
     GraphDataLoader,
     GraphDataset,
     GraphSample,
+    TensorGraph,
     batch_graph_samples,
 )
 from .metadata import (
@@ -26,6 +27,7 @@ __all__ = [
     "EventMetadata",
     "BatchMetadata",
     "GraphSample",
+    "TensorGraph",
     "GraphBatch",
     "GraphDataset",
     "GraphDataLoader",

src/gnn4colliders/data/graph_dataset.py

@@ -52,18 +52,102 @@ class GraphBatch:
         )
 
 
-def batch_graph_samples(samples: Sequence[GraphSample]) -> GraphBatch:
-    """Batch samples in the supplied order using DGL's native batching."""
+@dataclass(frozen=True)
+class TensorGraph:
+    """Batched graph tensors usable on CPU, CUDA, and Apple's MPS backend."""
+
+    node_features: torch.Tensor
+    edge_features: torch.Tensor
+    edge_src: torch.Tensor
+    edge_dst: torch.Tensor
+    node_batch: torch.Tensor
+    edge_batch: torch.Tensor
+    node_counts: torch.Tensor
+    edge_counts: torch.Tensor
+
+    def to(self, device: torch.device | str) -> "TensorGraph":
+        return TensorGraph(
+            *(
+                value.to(device)
+                for value in (
+                    self.node_features,
+                    self.edge_features,
+                    self.edge_src,
+                    self.edge_dst,
+                    self.node_batch,
+                    self.edge_batch,
+                    self.node_counts,
+                    self.edge_counts,
+                )
+            )
+        )
+
+
+def _tensor_graphs(samples: Sequence[GraphSample]) -> TensorGraph:
+    node_features: list[torch.Tensor] = []
+    edge_features: list[torch.Tensor] = []
+    edge_src: list[torch.Tensor] = []
+    edge_dst: list[torch.Tensor] = []
+    node_counts: list[int] = []
+    edge_counts: list[int] = []
+    offset = 0
+    for sample in samples:
+        graph = sample.graph
+        if isinstance(graph, TensorGraph):
+            if graph.node_counts.numel() != 1:
+                raise ValueError("each GraphSample must contain one graph")
+            nodes, edges = graph.node_features, graph.edge_features
+            src, dst = graph.edge_src, graph.edge_dst
+        else:
+            try:
+                nodes = graph.ndata["features"]
+                edges = graph.edata["features"]
+                src, dst = graph.edges(order="eid")
+            except AttributeError as error:
+                raise TypeError(
+                    "GraphSample graph must be DGL-like or TensorGraph"
+                ) from error
+        node_count, edge_count = int(nodes.shape[0]), int(edges.shape[0])
+        node_features.append(nodes)
+        edge_features.append(edges)
+        edge_src.append(src.to(dtype=torch.long) + offset)
+        edge_dst.append(dst.to(dtype=torch.long) + offset)
+        node_counts.append(node_count)
+        edge_counts.append(edge_count)
+        offset += node_count
+    node_count_tensor = torch.tensor(node_counts, dtype=torch.long)
+    edge_count_tensor = torch.tensor(edge_counts, dtype=torch.long)
+    return TensorGraph(
+        node_features=torch.cat(node_features),
+        edge_features=torch.cat(edge_features),
+        edge_src=torch.cat(edge_src),
+        edge_dst=torch.cat(edge_dst),
+        node_batch=torch.repeat_interleave(
+            torch.arange(len(samples)), node_count_tensor
+        ),
+        edge_batch=torch.repeat_interleave(
+            torch.arange(len(samples)), edge_count_tensor
+        ),
+        node_counts=node_count_tensor,
+        edge_counts=edge_count_tensor,
+    )
+
+
+def batch_graph_samples(
+    samples: Sequence[GraphSample], *, tensor_graphs: bool = False
+) -> GraphBatch:
+    """Batch samples into either DGL or portable tensor graph batches."""
 
     if not samples:
         raise ValueError("cannot batch an empty sequence")
-    try:
-        import dgl
-    except ImportError as error:  # pragma: no cover - optional dependency
-        raise ImportError(
-            "batch_graph_samples requires the 'root-gnn' extra"
-        ) from error
-    graphs = dgl.batch([sample.graph for sample in samples])
+    if tensor_graphs:
+        graphs = _tensor_graphs(samples)
+    else:
+        try:
+            import dgl
+        except ImportError as error:  # pragma: no cover - optional dependency
+            raise ImportError("DGL batching requires the 'root-gnn' extra") from error
+        graphs = dgl.batch([sample.graph for sample in samples])
     labels = torch.stack([torch.as_tensor(sample.label) for sample in samples])
     globals_ = [sample.global_features for sample in samples]
     if any(value is None for value in globals_):
@@ -97,6 +181,7 @@ class GraphDataLoader:
         shuffle: bool = False,
         drop_last: bool = False,
         seed: int = 12345,
+        tensor_graphs: bool = False,
     ) -> None:
         if batch_size < 1:
             raise ValueError("batch_size must be positive")
@@ -105,6 +190,7 @@ class GraphDataLoader:
         self.shuffle = shuffle
         self.drop_last = drop_last
         self.seed = seed
+        self.tensor_graphs = tensor_graphs
         self.epoch = 0
 
     def set_epoch(self, epoch: int) -> None:
@@ -126,7 +212,7 @@ class GraphDataLoader:
                 self.dataset[int(index)]
                 for index in indices[start : start + self.batch_size]
             ]
-            yield batch_graph_samples(selected)
+            yield batch_graph_samples(selected, tensor_graphs=self.tensor_graphs)
 
 
 class DistributedGraphDataLoader(GraphDataLoader):
@@ -146,9 +232,15 @@ class DistributedGraphDataLoader(GraphDataLoader):
         shuffle=False,
         drop_last=False,
         seed=12345,
+        tensor_graphs=False,
     ):
         super().__init__(
-            dataset, batch_size, shuffle=shuffle, drop_last=drop_last, seed=seed
+            dataset,
+            batch_size,
+            shuffle=shuffle,
+            drop_last=drop_last,
+            seed=seed,
+            tensor_graphs=tensor_graphs,
         )
         from gnn4colliders.distributed import DistributedIndices
 
@@ -181,7 +273,8 @@ class DistributedGraphDataLoader(GraphDataLoader):
                 [
                     self.dataset[index]
                     for index in indices[start : start + self.batch_size]
-                ]
+                ],
+                tensor_graphs=self.tensor_graphs,
             )
 
 

src/gnn4colliders/models/root_gnn/__init__.py

@@ -1,6 +1,12 @@
 """ROOT-GNN model components."""
 
 from .edge_network import EdgeNetwork
+from .tensor_edge_network import TensorEdgeNetwork
 from .transfer import FineTunedEdgeNetwork, load_legacy_edge_network_state_dict
 
-__all__ = ["EdgeNetwork", "FineTunedEdgeNetwork", "load_legacy_edge_network_state_dict"]
+__all__ = [
+    "EdgeNetwork",
+    "TensorEdgeNetwork",
+    "FineTunedEdgeNetwork",
+    "load_legacy_edge_network_state_dict",
+]

src/gnn4colliders/models/root_gnn/tensor_edge_network.py

@@ -0,0 +1,173 @@
+"""Portable tensor-only ROOT-GNN message passing."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch import nn
+
+from gnn4colliders.data import TensorGraph
+
+from .blocks import make_mlp
+
+
+def _sum_by_group(
+    values: torch.Tensor, groups: torch.Tensor, count: int
+) -> torch.Tensor:
+    result = values.new_zeros((count, values.shape[1]))
+    return result.index_add_(0, groups, values)
+
+
+def _mean_by_group(
+    values: torch.Tensor, groups: torch.Tensor, count: int
+) -> torch.Tensor:
+    sums = _sum_by_group(values, groups, count)
+    sizes = _sum_by_group(values.new_ones((values.shape[0], 1)), groups, count)
+    return sums / sizes.clamp_min(1)
+
+
+class TensorEdgeNetwork(nn.Module):
+    """ROOT-GNN using only PyTorch tensors, including on MPS and CUDA.
+
+    Parameter names and update equations intentionally match :class:`EdgeNetwork`
+    so historical checkpoints load without conversion.
+    """
+
+    def __init__(
+        self,
+        sample_graph: TensorGraph,
+        sample_global: torch.Tensor | None,
+        hid_size: int,
+        out_size: int,
+        n_layers: int,
+        n_proc_steps: int,
+        dropout: float = 0.0,
+        **_: Any,
+    ) -> None:
+        super().__init__()
+        if not isinstance(sample_graph, TensorGraph):
+            raise TypeError("TensorEdgeNetwork requires a TensorGraph sample")
+        if n_proc_steps < 0:
+            raise ValueError("n_proc_steps must be non-negative")
+        global_width = 0 if sample_global is None else sample_global.shape[1]
+        self.has_global = global_width != 0
+        encoder_global_width = global_width or 1
+        self.hid_size = hid_size
+        self.n_layers = n_layers
+        self.n_proc_steps = n_proc_steps
+        self.node_feature_size = int(sample_graph.node_features.shape[1])
+        self.edge_feature_size = int(sample_graph.edge_features.shape[1])
+        self.global_feature_size = int(global_width)
+        self.dropout = float(dropout)
+        self.node_encoder = make_mlp(
+            self.node_feature_size, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.edge_encoder = make_mlp(
+            self.edge_feature_size, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.global_encoder = make_mlp(
+            encoder_global_width, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.node_update = make_mlp(
+            3 * hid_size, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.edge_update = make_mlp(
+            4 * hid_size, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.global_update = make_mlp(
+            3 * hid_size, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.global_decoder = make_mlp(
+            hid_size, hid_size, hid_size, n_layers, dropout=dropout
+        )
+        self.classifier = nn.Linear(hid_size, out_size)
+
+    @property
+    def classify(self) -> nn.Linear:
+        return self.classifier
+
+    def checkpoint_config(self) -> dict[str, Any]:
+        return {
+            "family": "root_gnn",
+            "class": type(self).__name__,
+            "backend": "torch",
+            "node_feature_size": self.node_feature_size,
+            "edge_feature_size": self.edge_feature_size,
+            "global_feature_size": self.global_feature_size,
+            "hid_size": self.hid_size,
+            "out_size": int(self.classifier.out_features),
+            "n_layers": self.n_layers,
+            "n_proc_steps": self.n_proc_steps,
+            "dropout": self.dropout,
+        }
+
+    def _globals(
+        self, graph: TensorGraph, global_features: torch.Tensor | None
+    ) -> torch.Tensor:
+        if not self.has_global:
+            return graph.node_counts[:, None].to(
+                device=graph.node_features.device, dtype=torch.float32
+            )
+        if global_features is None:
+            raise ValueError("global_features are required for this model")
+        if global_features.ndim == 1:
+            if graph.node_counts.numel() != 1:
+                raise ValueError("one-dimensional global_features require one graph")
+            return global_features.unsqueeze(0)
+        return global_features
+
+    def forward_features(
+        self, graph: TensorGraph, global_features: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        if hasattr(graph, "graph") and hasattr(graph, "global_features"):
+            global_features, graph = graph.global_features, graph.graph
+        if not isinstance(graph, TensorGraph):
+            raise TypeError("TensorEdgeNetwork requires TensorGraph inputs")
+        node_h = self.node_encoder(graph.node_features)
+        edge_h = self.edge_encoder(graph.edge_features)
+        global_h = self.global_encoder(self._globals(graph, global_features))
+        graph_count = int(graph.node_counts.numel())
+        for _ in range(self.n_proc_steps):
+            edge_h = self.edge_update(
+                torch.cat(
+                    (
+                        edge_h,
+                        node_h[graph.edge_src],
+                        node_h[graph.edge_dst],
+                        global_h[graph.edge_batch],
+                    ),
+                    dim=1,
+                )
+            )
+            node_h = self.node_update(
+                torch.cat(
+                    (
+                        node_h,
+                        _sum_by_group(edge_h, graph.edge_dst, node_h.shape[0]),
+                        global_h[graph.node_batch],
+                    ),
+                    dim=1,
+                )
+            )
+            global_h = self.global_update(
+                torch.cat(
+                    (
+                        global_h,
+                        _mean_by_group(node_h, graph.node_batch, graph_count),
+                        _mean_by_group(edge_h, graph.edge_batch, graph_count),
+                    ),
+                    dim=1,
+                )
+            )
+        return self.global_decoder(global_h)
+
+    def forward(
+        self, graph: TensorGraph, global_features: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        return self.classifier(self.forward_features(graph, global_features))
+
+    def representation(
+        self, graph: TensorGraph, global_features: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        return self.forward_features(graph, global_features)

tests/parity/conftest.py

@@ -68,6 +68,11 @@ def legacy_dataset_module_without_dgl():
     try:
         import dgl  # noqa: F401
     except ImportError:
+        missing = object()
+        originals = {
+            name: sys.modules.get(name, missing)
+            for name in ("dgl", "dgl.data", "matplotlib", "matplotlib.pyplot")
+        }
         dgl_module = types.ModuleType("dgl")
         dgl_data_module = types.ModuleType("dgl.data")
         dgl_data_module.DGLDataset = type("DGLDataset", (), {})
@@ -79,6 +84,15 @@ def legacy_dataset_module_without_dgl():
         matplotlib_module.pyplot = matplotlib_pyplot_module
         sys.modules["matplotlib"] = matplotlib_module
         sys.modules["matplotlib.pyplot"] = matplotlib_pyplot_module
+        try:
+            from root_gnn_base import dataset
+        finally:
+            for name, original in originals.items():
+                if original is missing:
+                    sys.modules.pop(name, None)
+                else:
+                    sys.modules[name] = original
+        return dataset
     from root_gnn_base import dataset
 
     return dataset

tests/unit/graphs/test_dgl_graph.py

@@ -1,7 +1,10 @@
+import pytest
 import torch
 
 from gnn4colliders.graphs import build_dgl_graph
 
+pytest.importorskip("dgl")
+
 
 def test_build_dgl_graph_attaches_compatible_features():
     node_features = torch.tensor(

tests/unit/models/root_gnn/test_tensor_edge_network.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from gnn4colliders.data import EventMetadata, GraphDataset, GraphSample
+from gnn4colliders.data.graph_dataset import GraphDataLoader
+from gnn4colliders.graphs import build_dgl_graph
+from gnn4colliders.models.root_gnn import EdgeNetwork, TensorEdgeNetwork
+
+pytest.importorskip("dgl")
+
+
+def _batch():
+    samples = [
+        GraphSample(
+            build_dgl_graph(nodes),
+            torch.tensor(index),
+            torch.tensor([1.0, 2.0]),
+            EventMetadata(0, 1.0, str(index)),
+        )
+        for index, nodes in enumerate(
+            (
+                torch.tensor([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]]),
+                torch.tensor([[4.0, 0.1, 2.8]]),
+            )
+        )
+    ]
+    dgl_batch = next(iter(GraphDataLoader(GraphDataset(samples), 2)))
+    tensor_batch = next(
+        iter(GraphDataLoader(GraphDataset(samples), 2, tensor_graphs=True))
+    )
+    return dgl_batch, tensor_batch
+
+
+def test_tensor_backend_matches_dgl_for_fixed_weights():
+    dgl_batch, tensor_batch = _batch()
+    dgl_model = EdgeNetwork(
+        dgl_batch.graph, dgl_batch.global_features, 8, 3, 2, 2
+    ).eval()
+    tensor_model = TensorEdgeNetwork(
+        tensor_batch.graph, tensor_batch.global_features, 8, 3, 2, 2
+    ).eval()
+    tensor_model.load_state_dict(dgl_model.state_dict())
+    with torch.no_grad():
+        expected = dgl_model(dgl_batch.graph, dgl_batch.global_features)
+        actual = tensor_model(tensor_batch.graph, tensor_batch.global_features)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    dgl_model(dgl_batch.graph, dgl_batch.global_features).sum().backward()
+    tensor_model(tensor_batch.graph, tensor_batch.global_features).sum().backward()
+    for expected_parameter, actual_parameter in zip(
+        dgl_model.parameters(), tensor_model.parameters(), strict=True
+    ):
+        torch.testing.assert_close(
+            actual_parameter.grad, expected_parameter.grad, rtol=0, atol=0
+        )
+
+
+@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS unavailable")
+def test_tensor_backend_runs_forward_and_backward_on_mps():
+    _, batch = _batch()
+    model = TensorEdgeNetwork(batch.graph, batch.global_features, 8, 3, 2, 1).to("mps")
+    moved = batch.to("mps")
+    logits = model(moved.graph, moved.global_features)
+    logits.square().mean().backward()
+    assert logits.device.type == "mps"
+    assert all(parameter.grad is not None for parameter in model.parameters())

tests/unit/training/test_trainer.py

@@ -90,3 +90,16 @@ def test_graph_batch_to_cuda_preserves_named_metadata():
     assert moved.labels.device.type == "cuda"
     assert moved.metadata.weight.device.type == "cuda"
     assert moved.metadata.sample_id == ("event:0",)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS unavailable")
+def test_trainer_runs_optimizer_step_on_mps():
+    model = TinyModel()
+    optimizer = build_optimizer(model, learning_rate=0.01)
+    trainer = Trainer(model, BinaryClassificationTask(), optimizer, device="mps")
+    before = trainer.model.linear.weight.detach().clone()
+    result = trainer.train_batch(_batch([1, 2], [0, 1]))
+    assert torch.isfinite(torch.tensor(result.loss))
+    assert trainer.model.linear.weight.device.type == "mps"
+    assert not torch.equal(before, trainer.model.linear.weight.detach())

tests/unit/validation/test_artifacts.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import numpy as np
-
 from validation.artifacts import ValidationArtifact, load_artifact, save_artifact
 from validation.compare import compare_artifacts