fix(pattern): skip ONNX validation when external data is unloaded

src/winml/modelkit/onnx/__init__.py

@@ -20,7 +20,7 @@
 from .metadata import capture_metadata, restore_metadata
 from .persistence import cleanup_onnx, load_onnx, save_onnx
 from .shape import infer_onnx_shapes, infer_shapes
-from .utils import EXTERNAL_DATA_THRESHOLD, check_onnx_model, get_model_size
+from .utils import EXTERNAL_DATA_THRESHOLD, check_onnx_model, get_model_size, has_unloaded_external_data
 
 
 __all__ = [
@@ -36,6 +36,7 @@
     "generate_inputs_from_onnx",
     "get_io_config",
     "get_model_size",
+    "has_unloaded_external_data",
     "infer_onnx_shapes",
     "infer_shapes",
     "is_compiled_onnx",

src/winml/modelkit/onnx/utils.py

@@ -13,6 +13,20 @@
 EXTERNAL_DATA_THRESHOLD = 100 * 1024 * 1024  # 100 MiB
 
 
+def has_unloaded_external_data(model: onnx.ModelProto) -> bool:
+    """Return True if the model contains external-data tensors whose bytes are not in memory.
+
+    When a model is loaded with ``load_external_data=False``, tensors keep their
+    ``data_location == EXTERNAL`` annotation but ``raw_data`` stays empty.
+    The ONNX checker needs the sidecar ``.data`` file on disk to validate those
+    tensors, which may not be available in the current working directory.
+    """
+    return any(
+        t.data_location == onnx.TensorProto.EXTERNAL and not t.raw_data
+        for t in _get_all_tensors(model)
+    )
+
+
 def get_model_size(model: onnx.ModelProto) -> int:
     """Calculate the total size of an ONNX model in bytes.
 
@@ -35,11 +49,22 @@ def check_onnx_model(
     full_check: bool = False,
     skip_opset_compatibility_check: bool = False,
     check_custom_domain: bool = False,
+    skip_if_unloaded_external_data: bool = False,
 ) -> None:
     """Same as ``onnx.checker.check_model``, but handles >2GiB models.
 
     Uses a temp file on disk for large models.
+
+    Args:
+        skip_if_unloaded_external_data: When True, skip validation if the model
+            has tensors with ``data_location == EXTERNAL`` but no ``raw_data``
+            (i.e. loaded with ``load_external_data=False``).  The ONNX checker
+            needs the sidecar ``.data`` file on disk to validate those tensors,
+            which is not available in that case.
     """
+    if skip_if_unloaded_external_data and has_unloaded_external_data(model):
+        return
+
     tmp_dir = None
 
     if get_model_size(model) >= EXTERNAL_DATA_THRESHOLD:

src/winml/modelkit/pattern/base.py

@@ -1238,6 +1238,12 @@ def __init__(self, onnx_model: ModelProto, raise_on_invalid_model: bool = True)
         # Maps tensor name -> numpy array value (for constants/initializers)
         self.tensor_values: dict[str, np.ndarray] = {}
 
+        # Names of initializers whose external data was not loaded into memory.
+        # Constant-value constraints are skipped for these tensors so that
+        # topology-based pattern matching still works on models loaded with
+        # load_external_data=False.
+        self._external_unloaded_names: set[str] = set()
+
         # Maps tensor name -> shape tuple
         self.tensor_shapes: dict[str, tuple] = {}
 
@@ -1255,7 +1261,7 @@ def __init__(self, onnx_model: ModelProto, raise_on_invalid_model: bool = True)
 
         if raise_on_invalid_model:
             try:
-                check_onnx_model(self.model)
+                check_onnx_model(self.model, skip_if_unloaded_external_data=True)
             except onnx.checker.ValidationError as e:
                 raise InvalidPatternMatcherModelError(
                     f"Model failed ONNX validation: {e}",
@@ -1320,6 +1326,11 @@ def _build_lookups(self) -> None:
         for initializer in self.graph.initializer:
             self.producer_lookup.setdefault(initializer.name, (initializer.name, 0, "Initializer"))
             if initializer.name:
+                if initializer.data_location == onnx.TensorProto.EXTERNAL and not initializer.raw_data:
+                    # External data not loaded; record the name so constant-value
+                    # checks can be skipped rather than returning a false failure.
+                    self._external_unloaded_names.add(initializer.name)
+                    continue
                 self.tensor_values[initializer.name] = numpy_helper.to_array(initializer)
 
         for node_idx, node in enumerate(self.graph.node):
@@ -1519,6 +1530,10 @@ def _check_constant_constraints(
 
             # Get tensor value
             if input_tensor_name not in self.tensor_values:
+                # If the initializer's external data was not loaded, skip the value
+                # check: topology matched, but we can't verify the scalar constant.
+                if input_tensor_name in self._external_unloaded_names:
+                    continue
                 return False
             actual_value = self.tensor_values[input_tensor_name]
 

tests/unit/analyze/pattern/test_pattern_matching.py

@@ -7,9 +7,12 @@
 
 from pathlib import Path
 
+import numpy as np
 import onnx
 import onnx.helper as oh
+import onnx.numpy_helper as onph
 import pytest
+from onnx import TensorProto
 
 from winml.modelkit.pattern import (
     Gelu2Pattern,
@@ -328,3 +331,78 @@ def test_raise_on_invalid_model_false_still_works(self):
         matcher.register_pattern(Gelu2Pattern())
         results = matcher.match()
         assert len(results) == 1
+
+
+def _make_gelu_model_with_external_initializers() -> onnx.ModelProto:
+    """Gelu2 model where weight initializers simulate load_external_data=False.
+
+    The graph topology is identical to a normal Gelu2 model; only the weight
+    tensors have data_location=EXTERNAL with empty raw_data, mimicking what
+    onnx.load(..., load_external_data=False) produces for large models.
+    """
+    X = oh.make_tensor_value_info("X", TensorProto.FLOAT, [1, 8])  # noqa: N806
+    Y = oh.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 8])  # noqa: N806
+
+    def _external_tensor(name: str, values: list, dtype=TensorProto.FLOAT, dims=None):
+        t = onph.from_array(
+            np.array(values, dtype=np.float32 if dtype == TensorProto.FLOAT else np.int64),
+            name=name,
+        )
+        t.data_location = TensorProto.EXTERNAL
+        t.ClearField("raw_data")
+        entry = t.external_data.add()
+        entry.key = "location"
+        entry.value = "model.onnx.data"
+        return t
+
+    sqrt2 = _external_tensor("sqrt2", [1.4142135])
+    half = _external_tensor("half", [0.5])
+    one = _external_tensor("one_val", [1.0])
+
+    div = oh.make_node("Div", ["X", "sqrt2"], ["div_out"], name="div_node")
+    erf = oh.make_node("Erf", ["div_out"], ["erf_out"], name="erf_node")
+    add = oh.make_node("Add", ["erf_out", "one_val"], ["add_out"], name="add_node")
+    mul1 = oh.make_node("Mul", ["X", "add_out"], ["mul1_out"], name="mul1_node")
+    mul2 = oh.make_node("Mul", ["mul1_out", "half"], ["Y"], name="mul2_node")
+
+    graph = oh.make_graph(
+        [div, erf, add, mul1, mul2],
+        "gelu_external",
+        [X],
+        [Y],
+        initializer=[sqrt2, half, one],
+    )
+    return oh.make_model(graph, opset_imports=[oh.make_opsetid("", 13)])
+
+
+class TestPatternMatchingWithUnloadedExternalData:
+    """PatternMatcher must not fail when model initializers have unloaded external data.
+
+    This covers the regression where models loaded with load_external_data=False
+    triggered an ONNX ValidationError in check_onnx_model because the sidecar
+    .data file could not be found, causing pattern matching to be silently skipped.
+    """
+
+    def test_pattern_matcher_does_not_raise(self):
+        model = _make_gelu_model_with_external_initializers()
+        matcher = PatternMatcher(model)
+        assert matcher is not None
+
+    def test_gelu_pattern_still_matched_on_topology(self):
+        """Topology-based pattern matching succeeds even without tensor values."""
+        model = _make_gelu_model_with_external_initializers()
+        matcher = PatternMatcher(model)
+        matcher.register_pattern(Gelu2Pattern())
+        results = matcher.match()
+        assert len(results) == 1, f"Expected 1 Gelu2Pattern match, got {len(results)}"
+
+    def test_unloaded_initializers_absent_from_tensor_values(self):
+        """External-data initializers must not pollute tensor_values with zero arrays."""
+        model = _make_gelu_model_with_external_initializers()
+        matcher = PatternMatcher(model)
+        # None of the external initializer names should have values loaded.
+        external_names = {init.name for init in model.graph.initializer}
+        for name in external_names:
+            assert name not in matcher.tensor_values, (
+                f"Initializer '{name}' should not be in tensor_values when external data is unloaded"
+            )

tests/unit/onnx/test_onnx_utils.py

@@ -0,0 +1,82 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Tests for modelkit.onnx.utils — check_onnx_model and has_unloaded_external_data."""
+
+from __future__ import annotations
+
+import numpy as np
+import onnx
+import pytest
+from onnx import TensorProto, helper, numpy_helper
+
+from winml.modelkit.onnx import check_onnx_model, has_unloaded_external_data
+
+
+def _make_simple_model() -> onnx.ModelProto:
+    """Minimal valid ONNX model with an inline initializer."""
+    x_info = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4])
+    y_info = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4])
+    weight = numpy_helper.from_array(np.ones((4,), dtype=np.float32), name="W")
+    node = helper.make_node("Add", ["X", "W"], ["Y"], name="add")
+    graph = helper.make_graph([node], "g", [x_info], [y_info], initializer=[weight])
+    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+
+
+def _make_model_with_unloaded_external_data() -> onnx.ModelProto:
+    """Model whose initializer mimics load_external_data=False: data_location=EXTERNAL, raw_data empty."""
+    x_info = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4])
+    y_info = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4])
+
+    weight = numpy_helper.from_array(np.ones((4,), dtype=np.float32), name="W")
+    # Simulate what onnx.load(..., load_external_data=False) produces for large tensors.
+    weight.data_location = TensorProto.EXTERNAL
+    weight.ClearField("raw_data")
+    entry = weight.external_data.add()
+    entry.key = "location"
+    entry.value = "model.onnx.data"
+
+    node = helper.make_node("Add", ["X", "W"], ["Y"], name="add")
+    graph = helper.make_graph([node], "g", [x_info], [y_info], initializer=[weight])
+    return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+
+
+class TestHasUnloadedExternalData:
+    def test_inline_model_returns_false(self):
+        assert has_unloaded_external_data(_make_simple_model()) is False
+
+    def test_unloaded_external_tensor_returns_true(self):
+        assert has_unloaded_external_data(_make_model_with_unloaded_external_data()) is True
+
+    def test_loaded_external_data_returns_false(self):
+        """After raw_data is populated the tensor is considered loaded."""
+        model = _make_model_with_unloaded_external_data()
+        init = model.graph.initializer[0]
+        init.raw_data = np.ones((4,), dtype=np.float32).tobytes()
+        assert has_unloaded_external_data(model) is False
+
+
+class TestCheckOnnxModel:
+    def test_valid_inline_model_does_not_raise(self):
+        check_onnx_model(_make_simple_model())
+
+    def test_unloaded_external_data_does_not_raise_when_skip_enabled(self):
+        """check_onnx_model must not raise when skip_if_unloaded_external_data=True."""
+        check_onnx_model(
+            _make_model_with_unloaded_external_data(),
+            skip_if_unloaded_external_data=True,
+        )
+
+    def test_unloaded_external_data_raises_by_default(self):
+        """check_onnx_model raises by default when external data is missing on disk."""
+        with pytest.raises(onnx.checker.ValidationError):
+            check_onnx_model(_make_model_with_unloaded_external_data())
+
+    def test_invalid_model_raises(self):
+        """check_onnx_model should still raise for structurally invalid models."""
+        model = _make_simple_model()
+        # Corrupt the opset version to something invalid.
+        model.opset_import[0].version = 0
+        with pytest.raises(onnx.checker.ValidationError):
+            check_onnx_model(model)