Add CI coverage that pretrained model configs and QUESO weights still load

src/graphnet/models/gnn/dynedge.py

@@ -1,6 +1,6 @@
 """Implementation of the DynEdge GNN model architecture."""
 
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Sequence, Union
 
 import torch
 from torch import Tensor, LongTensor
@@ -28,7 +28,7 @@ def __init__(
         *,
         nb_neighbours: int = 8,
         features_subset: Optional[Union[List[int], slice]] = None,
-        dynedge_layer_sizes: Optional[List[Tuple[int, ...]]] = None,
+        dynedge_layer_sizes: Optional[List[Sequence[int]]] = None,
         post_processing_layer_sizes: Optional[List[int]] = None,
         readout_layer_sizes: Optional[List[int]] = None,
         global_pooling_schemes: Optional[Union[str, List[str]]] = None,
@@ -102,7 +102,10 @@ def __init__(
 
         assert isinstance(dynedge_layer_sizes, list)
         assert len(dynedge_layer_sizes)
-        assert all(isinstance(sizes, tuple) for sizes in dynedge_layer_sizes)
+        # YAML has no tuple type, so a serialized config reloads these inner
+        # size pairs as lists; coerce them back so a saved DynEdge config
+        # round-trips.
+        dynedge_layer_sizes = [tuple(sizes) for sizes in dynedge_layer_sizes]
         assert all(len(sizes) > 0 for sizes in dynedge_layer_sizes)
         assert all(
             all(size > 0 for size in sizes) for sizes in dynedge_layer_sizes

tests/models/test_pretrained.py

@@ -0,0 +1,65 @@
+"""Unit tests for the pretrained models shipped with GraphNeT."""
+
+import glob
+import os
+
+import pytest
+
+from graphnet.constants import PRETRAINED_MODEL_DIR
+from graphnet.models import Model
+from graphnet.utilities.config import ModelConfig
+
+
+def _config_paths() -> list:
+    """Return every pretrained model config shipped in the repository."""
+    return sorted(
+        glob.glob(
+            os.path.join(PRETRAINED_MODEL_DIR, "**", "*.yml"), recursive=True
+        )
+    )
+
+
+def _config_id(path: str) -> str:
+    """Return a readable test id relative to the pretrained model dir."""
+    return os.path.relpath(path, PRETRAINED_MODEL_DIR)
+
+
+def _config_paths_with_state_dict() -> list:
+    """Return pretrained configs shipped alongside a state dict."""
+    return [
+        path
+        for path in _config_paths()
+        if path.endswith("_config.yml")
+        and os.path.exists(path.replace("_config.yml", "_state_dict.pth"))
+    ]
+
+
+@pytest.mark.parametrize("config_path", _config_paths(), ids=_config_id)
+def test_pretrained_config_builds(config_path: str) -> None:
+    """Test that every shipped pretrained config constructs a model.
+
+    Guards the committed configs against silent rot when a constructor
+    argument or class name changes elsewhere in the library.
+    """
+    config = ModelConfig.load(config_path)
+    assert isinstance(config, ModelConfig)
+    model = Model.from_config(config, trust=True)
+    assert isinstance(model, Model)
+
+
+@pytest.mark.parametrize(
+    "config_path", _config_paths_with_state_dict(), ids=_config_id
+)
+def test_pretrained_state_dict_loads(config_path: str) -> None:
+    """Test that shipped weights load into their model without key mismatch.
+
+    Only applies to models whose state dict is committed next to the
+    config; a strict load verifies the weights and the current
+    architecture still agree exactly.
+    """
+    config = ModelConfig.load(config_path)
+    assert isinstance(config, ModelConfig)
+    model = Model.from_config(config, trust=True)
+
+    state_dict_path = config_path.replace("_config.yml", "_state_dict.pth")
+    model.load_state_dict(state_dict_path)