fix(export): align ONNX input names to forward() signature order

examples/recipes/google_owlvit-base-patch32/cpu/cpu/zero-shot-object-detection_fp16_config.json

@@ -0,0 +1,77 @@
+{
+  "export": {
+    "opset_version": 17,
+    "batch_size": 1,
+    "export_params": true,
+    "do_constant_folding": true,
+    "verbose": false,
+    "dynamo": false,
+    "enable_hierarchy_tags": true,
+    "clean_onnx": false,
+    "hierarchy_tag_format": "full",
+    "input_tensors": [
+      {
+        "name": "pixel_values",
+        "dtype": "float32",
+        "shape": [1, 3, 768, 768],
+        "value_range": [0, 1]
+      },
+      {
+        "name": "input_ids",
+        "dtype": "int32",
+        "shape": [1, 16],
+        "value_range": [0, 49408]
+      },
+      {
+        "name": "attention_mask",
+        "dtype": "int32",
+        "shape": [1, 16],
+        "value_range": [0, 2]
+      }
+    ],
+    "output_tensors": [
+      {
+        "name": "logits"
+      },
+      {
+        "name": "pred_boxes"
+      },
+      {
+        "name": "text_embeds"
+      },
+      {
+        "name": "image_embeds"
+      }
+    ]
+  },
+  "optim": {},
+  "quant": {
+    "mode": "fp16",
+    "samples": 10,
+    "calibration_method": "minmax",
+    "weight_type": "uint8",
+    "activation_type": "uint8",
+    "per_channel": false,
+    "symmetric": false,
+    "weight_symmetric": null,
+    "activation_symmetric": null,
+    "save_calibration": false,
+    "distribution": "uniform",
+    "seed": null,
+    "calibration_load_path": null,
+    "calibration_save_path": null,
+    "op_types_to_quantize": null,
+    "nodes_to_exclude": null,
+    "task": "zero-shot-object-detection",
+    "model_id": "google/owlvit-base-patch32",
+    "model_type": "owlvit",
+    "fp16_keep_io_types": true,
+    "fp16_op_block_list": null
+  },
+  "compile": null,
+  "loader": {
+    "task": "zero-shot-object-detection",
+    "model_class": "AutoModelForZeroShotObjectDetection",
+    "model_type": "owlvit"
+  }
+}

examples/recipes/google_owlvit-base-patch32/cpu/cpu/zero-shot-object-detection_fp32_config.json

@@ -0,0 +1,55 @@
+{
+  "export": {
+    "opset_version": 17,
+    "batch_size": 1,
+    "export_params": true,
+    "do_constant_folding": true,
+    "verbose": false,
+    "dynamo": false,
+    "enable_hierarchy_tags": true,
+    "clean_onnx": false,
+    "hierarchy_tag_format": "full",
+    "input_tensors": [
+      {
+        "name": "pixel_values",
+        "dtype": "float32",
+        "shape": [1, 3, 768, 768],
+        "value_range": [0, 1]
+      },
+      {
+        "name": "input_ids",
+        "dtype": "int32",
+        "shape": [1, 16],
+        "value_range": [0, 49408]
+      },
+      {
+        "name": "attention_mask",
+        "dtype": "int32",
+        "shape": [1, 16],
+        "value_range": [0, 2]
+      }
+    ],
+    "output_tensors": [
+      {
+        "name": "logits"
+      },
+      {
+        "name": "pred_boxes"
+      },
+      {
+        "name": "text_embeds"
+      },
+      {
+        "name": "image_embeds"
+      }
+    ]
+  },
+  "optim": {},
+  "quant": null,
+  "compile": null,
+  "loader": {
+    "task": "zero-shot-object-detection",
+    "model_class": "AutoModelForZeroShotObjectDetection",
+    "model_type": "owlvit"
+  }
+}

src/winml/modelkit/export/htp/exporter.py

@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 import contextlib
+import inspect
 import logging
 import sys
 import time
@@ -46,6 +47,40 @@
 logger = logging.getLogger(__name__)
 
 
+def _reorder_inputs_by_forward_signature(
+    model: nn.Module,
+    input_names: list[str],
+    inputs: dict,
+) -> tuple[list[str], dict]:
+    """Reorder input_names and inputs dict to match model.forward() signature.
+
+    PyTorch's torch.onnx.export lays out graph inputs in forward() parameter
+    order regardless of dict key order, then assigns input_names positionally.
+    If input_names follow a different order (e.g. OnnxConfig.inputs declaration
+    order), labels get swapped.  This function aligns both to forward() order,
+    equivalent to Optimum's ``ordered_inputs(model)`` helper.
+
+    Falls back to the original order if inspection fails or if the input names
+    don't overlap with forward() parameters (e.g. wrapper models).
+    """
+    try:
+        forward_params = list(inspect.signature(model.forward).parameters.keys())
+    except (ValueError, TypeError):
+        return input_names, inputs
+
+    # Build ordered list: params that appear in input_names, in forward order
+    input_set = set(input_names)
+    ordered = [p for p in forward_params if p in input_set]
+
+    # If not all input_names matched (wrapper model, renamed params), keep original
+    if set(ordered) != input_set:
+        return input_names, inputs
+
+    # Rebuild inputs dict in forward order
+    ordered_inputs = {name: inputs[name] for name in ordered}
+    return ordered, ordered_inputs
+
+
 class HTPConfig:
     """Configuration constants for HTP Exporter."""
 
@@ -448,6 +483,17 @@ def _convert_model_to_onnx(
         # Input names from config, fallback to inputs dict keys
         input_names = export_config.get_input_names() or list(inputs.keys())
 
+        # Reorder input_names and inputs to match model.forward() signature.
+        # Optimum's OnnxConfig.inputs declares names in an arbitrary order that
+        # may differ from the model's forward() parameter order.  PyTorch's
+        # torch.onnx.export assigns input_names positionally to graph inputs
+        # which follow forward() signature order, so a mismatch causes labels
+        # to be swapped (e.g. pixel_values ↔ input_ids).  Aligning here is the
+        # equivalent of Optimum's ordered_inputs() helper.
+        input_names, inputs = _reorder_inputs_by_forward_signature(
+            model, input_names, inputs
+        )
+
         # Output names: infer from traced hierarchy, validate against config
         traced_outputs = self._hierarchy_builder.get_outputs() if self._hierarchy_builder else None
         inferred_names = infer_output_names(traced_outputs) if traced_outputs is not None else []

tests/unit/export/test_reorder_inputs.py

@@ -0,0 +1,95 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Tests for _reorder_inputs_by_forward_signature in htp/exporter.py."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from winml.modelkit.export.htp.exporter import _reorder_inputs_by_forward_signature
+
+
+class _MultiInputModel(nn.Module):
+    """Fake model whose forward() signature order differs from config order."""
+
+    def forward(
+        self, input_ids: torch.Tensor, pixel_values: torch.Tensor, attention_mask: torch.Tensor
+    ) -> torch.Tensor:
+        return input_ids + pixel_values + attention_mask
+
+
+class _SingleInputModel(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+
+class _WrapperModel(nn.Module):
+    """Model whose forward params don't match the config names at all."""
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states
+
+
+@pytest.mark.parametrize(
+    "config_order,expected_order",
+    [
+        # Config declares different order from forward signature
+        (
+            ["pixel_values", "input_ids", "attention_mask"],
+            ["input_ids", "pixel_values", "attention_mask"],
+        ),
+        # Config already matches forward signature
+        (
+            ["input_ids", "pixel_values", "attention_mask"],
+            ["input_ids", "pixel_values", "attention_mask"],
+        ),
+        # Subset of forward params (only 2 of 3)
+        (
+            ["attention_mask", "input_ids"],
+            ["input_ids", "attention_mask"],
+        ),
+    ],
+)
+def test_reorder_aligns_to_forward_signature(config_order, expected_order):
+    """Input names and dict are reordered to match model.forward() parameter order."""
+    model = _MultiInputModel()
+    inputs = {name: torch.zeros(1) for name in config_order}
+
+    ordered_names, ordered_inputs = _reorder_inputs_by_forward_signature(
+        model, config_order, inputs
+    )
+
+    assert ordered_names == expected_order
+    assert list(ordered_inputs.keys()) == expected_order
+
+
+def test_reorder_falls_back_when_names_dont_match():
+    """If config names don't overlap with forward params, original order is kept."""
+    model = _WrapperModel()
+    config_order = ["pixel_values", "input_ids"]
+    inputs = {name: torch.zeros(1) for name in config_order}
+
+    ordered_names, ordered_inputs = _reorder_inputs_by_forward_signature(
+        model, config_order, inputs
+    )
+
+    # Should be unchanged
+    assert ordered_names == config_order
+    assert list(ordered_inputs.keys()) == config_order
+
+
+def test_reorder_single_input_is_noop():
+    """Single-input model is trivially correct regardless of order."""
+    model = _SingleInputModel()
+    config_order = ["x"]
+    inputs = {"x": torch.zeros(1)}
+
+    ordered_names, _ordered_inputs = _reorder_inputs_by_forward_signature(
+        model, config_order, inputs
+    )
+
+    assert ordered_names == ["x"]