test(e2e): add dynamic-axes e2e tests for export, config, build, perf

tests/e2e/test_build_e2e.py

@@ -30,6 +30,7 @@
 from typing import TYPE_CHECKING
 from unittest.mock import MagicMock, patch
 
+import onnx
 import pytest
 from click.testing import CliRunner
 
@@ -293,3 +294,67 @@ def test_onnx_passthrough_no_optimize(self, tmp_path: Path, onnx_model_path: Pat
         )
         assert result.exit_code == 0, f"build failed (exit {result.exit_code}):\n{result.output}"
         assert output_dir.exists()
+
+
+# ===========================================================================
+# Dynamic axes: --dynamic-axes survives the full build pipeline.
+# ===========================================================================
+
+
+@pytest.mark.slow
+@pytest.mark.network
+class TestBuildDynamicAxes:
+    """``--dynamic-axes`` is applied during export and survives optimization."""
+
+    def test_dynamic_batch_survives_build(self, tmp_path: Path):
+        """A dynamic batch axis stays symbolic in the built ONNX artifact."""
+        config_path = _generate_config_file(
+            tmp_path,
+            "microsoft/resnet-50",
+            task="image-classification",
+        )
+        axes = tmp_path / "axes.json"
+        axes.write_text(json.dumps({"pixel_values": {"0": "batch"}}))
+        output_dir = tmp_path / "output"
+
+        result = CliRunner().invoke(
+            build,
+            [
+                "-c",
+                config_path,
+                "-m",
+                "microsoft/resnet-50",
+                "-o",
+                str(output_dir),
+                "--no-quant",
+                "--no-compile",
+                "--no-analyze",
+                "--dynamic-axes",
+                str(axes),
+            ],
+            obj={"debug": True},
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0, f"build failed (exit {result.exit_code}):\n{result.output}"
+
+        onnx_files = list(output_dir.rglob("*.onnx"))
+        assert onnx_files, (
+            f"No ONNX files found in {output_dir}. Contents: "
+            f"{[str(p) for p in output_dir.rglob('*')]}"
+        )
+
+        # Locate the graph carrying the model input (skip any auxiliary files).
+        model = None
+        for path in onnx_files:
+            candidate = onnx.load(str(path))
+            if any(i.name == "pixel_values" for i in candidate.graph.input):
+                model = candidate
+                break
+        assert model is not None, (
+            f"no ONNX with a 'pixel_values' input among {[str(p) for p in onnx_files]}"
+        )
+
+        pixel_values = next(i for i in model.graph.input if i.name == "pixel_values")
+        dims = pixel_values.type.tensor_type.shape.dim
+        assert dims[0].dim_param == "batch", f"expected a symbolic batch dim, got {list(dims)}"
+        assert [d.dim_value for d in dims[1:]] == [3, 224, 224]

tests/e2e/test_config_e2e.py

@@ -539,3 +539,38 @@ def test_explicit_fp16_still_triggers_quant(self) -> None:
         quant = data.get("quant")
         assert quant is not None, "Explicit --precision fp16 should produce a quant config"
         assert quant.get("mode") == "fp16"
+
+
+# ===========================================================================
+# Dynamic axes: --dynamic-axes
+# ===========================================================================
+
+
+class TestConfigDynamicAxes:
+    """``--dynamic-axes`` is recorded in the generated ``export`` config section.
+
+    JSON serialization keeps axis keys as strings, so the round-tripped mapping
+    is ``{"pixel_values": {"0": "batch"}}``.
+    """
+
+    MODEL = "microsoft/resnet-50"
+
+    def test_dynamic_axes_recorded(self, tmp_path: Path) -> None:
+        axes = tmp_path / "axes.json"
+        axes.write_text(json.dumps({"pixel_values": {"0": "batch"}}))
+        data = _run_config("-m", self.MODEL, "--dynamic-axes", str(axes))
+        _assert_hf_config_structure(data)
+        assert data["export"]["dynamic_axes"] == {"pixel_values": {"0": "batch"}}
+
+    def test_dynamic_axes_absent_by_default(self) -> None:
+        data = _run_config("-m", self.MODEL)
+        _assert_hf_config_structure(data)
+        assert data["export"].get("dynamic_axes") is None
+
+    def test_multiple_axes_recorded(self, tmp_path: Path) -> None:
+        mapping = {"pixel_values": {"0": "batch", "2": "height", "3": "width"}}
+        axes = tmp_path / "axes.json"
+        axes.write_text(json.dumps(mapping))
+        data = _run_config("-m", self.MODEL, "--dynamic-axes", str(axes))
+        _assert_hf_config_structure(data)
+        assert data["export"]["dynamic_axes"] == mapping

tests/e2e/test_export_e2e.py

@@ -186,6 +186,30 @@ def _write_json(path: Path, payload: dict) -> Path:
     return path
 
 
+def _symbolic_dims(tensors, name: str) -> dict[int, str]:
+    """Return ``{axis: dim_param}`` for the symbolic dims of the named tensor."""
+    for tensor in tensors:
+        if tensor.name == name:
+            return {
+                i: d.dim_param
+                for i, d in enumerate(tensor.type.tensor_type.shape.dim)
+                if d.dim_param
+            }
+    raise AssertionError(f"tensor {name!r} not found in {[t.name for t in tensors]}")
+
+
+def _static_dims(tensors, name: str) -> dict[int, int]:
+    """Return ``{axis: dim_value}`` for the static (non-symbolic) dims of the tensor."""
+    for tensor in tensors:
+        if tensor.name == name:
+            return {
+                i: d.dim_value
+                for i, d in enumerate(tensor.type.tensor_type.shape.dim)
+                if not d.dim_param
+            }
+    raise AssertionError(f"tensor {name!r} not found in {[t.name for t in tensors]}")
+
+
 # ===========================================================================
 # --help
 # ===========================================================================
@@ -220,17 +244,13 @@ def test_minimal_resnet50(self, tmp_path: Path):
 
 
 class TestExportDinoV2:
-
     MODEL = "facebook/dinov2-base"
 
     def test_image_feature_extraction(self, tmp_path: Path):
         """``-t image-feature-extraction`` must produce a valid ONNX export."""
         onnx_path = tmp_path / "model.onnx"
-        result = _invoke(["-m", self.MODEL, "-o", str(onnx_path),
-                          "-t", "image-feature-extraction"])
-        assert result.exit_code == 0, (
-            f"export failed (exit {result.exit_code}):\n{result.output}"
-        )
+        result = _invoke(["-m", self.MODEL, "-o", str(onnx_path), "-t", "image-feature-extraction"])
+        assert result.exit_code == 0, f"export failed (exit {result.exit_code}):\n{result.output}"
         assert onnx_path.exists(), f"ONNX model not found at {onnx_path}"
 
         model = onnx.load(str(onnx_path))
@@ -430,3 +450,56 @@ def test_build_config_with_dynamo(self, tmp_path: Path):
             f"expected opset 18 with -c + --dynamo, got {_opset_version(model)}"
         )
         _assert_some_node_has(model, "pkg.onnxscript.rewriter.rule_name")
+
+
+# ===========================================================================
+# Dynamic axes: --dynamic-axes
+# ===========================================================================
+
+
+class TestExportDynamicAxes:
+    """``--dynamic-axes`` marks the named tensor axes symbolic in the ONNX graph.
+
+    ResNet-50 has a single input ``pixel_values`` of shape [1, 3, 224, 224] and
+    a ``logits`` output. Making axis 0 dynamic turns the static batch dim into a
+    symbolic ``dim_param`` that also propagates to the output.
+    """
+
+    def test_batch_axis_symbolic(self, tmp_path: Path):
+        onnx_path = tmp_path / "model.onnx"
+        axes = _write_json(tmp_path / "axes.json", {"pixel_values": {"0": "batch"}})
+        model = _assert_succeeds(_happy_args(onnx_path, "--dynamic-axes", str(axes)), onnx_path)
+        # Axis 0 of pixel_values becomes symbolic; channel/spatial dims stay static.
+        assert _symbolic_dims(model.graph.input, "pixel_values") == {0: "batch"}
+        assert _static_dims(model.graph.input, "pixel_values") == {1: 3, 2: 224, 3: 224}
+        # The symbolic batch dim propagates to the logits output.
+        assert _symbolic_dims(model.graph.output, "logits") == {0: "batch"}
+
+    def test_static_batch_without_flag(self, tmp_path: Path):
+        # Baseline contrast: absent the flag, every input dim is a fixed integer.
+        onnx_path = tmp_path / "model.onnx"
+        model = _assert_succeeds(_happy_args(onnx_path), onnx_path)
+        assert _symbolic_dims(model.graph.input, "pixel_values") == {}
+        assert _input_shape_dims(model) == [1, 3, 224, 224]
+
+    def test_multiple_axes_symbolic(self, tmp_path: Path):
+        onnx_path = tmp_path / "model.onnx"
+        axes = _write_json(
+            tmp_path / "axes.json",
+            {"pixel_values": {"0": "batch", "2": "height", "3": "width"}},
+        )
+        model = _assert_succeeds(_happy_args(onnx_path, "--dynamic-axes", str(axes)), onnx_path)
+        assert _symbolic_dims(model.graph.input, "pixel_values") == {
+            0: "batch",
+            2: "height",
+            3: "width",
+        }
+        # Only the channel dim remains static.
+        assert _static_dims(model.graph.input, "pixel_values") == {1: 3}
+
+    def test_invalid_dynamic_axes_fails(self, tmp_path: Path):
+        # An empty symbolic dim name is rejected by WinMLExportConfig validation,
+        # so the command must fail cleanly without writing an ONNX file.
+        onnx_path = tmp_path / "model.onnx"
+        bad = _write_json(tmp_path / "axes.json", {"pixel_values": {"0": ""}})
+        _assert_fails(_happy_args(onnx_path, "--dynamic-axes", str(bad)), onnx_path)

tests/e2e/test_perf_e2e.py

@@ -35,6 +35,7 @@
 import sys
 from pathlib import Path
 
+import onnx
 import pytest
 from click.testing import CliRunner
 
@@ -741,6 +742,67 @@ def test_module_invalid_lists_available(self, tmp_path: Path):
         assert not output_file.exists(), "Output file should not be written on failure"
 
 
+# ===========================================================================
+# Dynamic axes: --dynamic-axes re-exports with a symbolic batch dim.
+# ===========================================================================
+
+
+class TestPerfDynamicAxes:
+    """``--dynamic-axes`` feeds the HF export so the benchmarked model is dynamic.
+
+    ``--ignore-cache`` forces a fresh build in a throwaway folder so the cached
+    static export isn't reused, and ``--no-skip-build`` guarantees the export
+    actually runs. The benchmarked ResNet-50 then exposes a dynamic (``None``)
+    batch axis instead of the static ``1``.
+    """
+
+    def test_dynamic_axes_cpu(self, tmp_path: Path):
+        axes = tmp_path / "axes.json"
+        axes.write_text(json.dumps({"pixel_values": {"0": "batch"}}))
+        output_file = tmp_path / "perf_dynamic_axes.json"
+
+        result = CliRunner().invoke(
+            perf,
+            [
+                "-m",
+                "microsoft/resnet-50",
+                "--iterations",
+                "3",
+                "--warmup",
+                "1",
+                "-o",
+                str(output_file),
+                "--device",
+                "cpu",
+                "--dynamic-axes",
+                str(axes),
+                "--ignore-cache",
+                "--no-skip-build",
+                "--no-optimize",
+                "--no-quant",
+                "--no-memory",
+            ],
+            obj={},
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0, f"perf failed (exit {result.exit_code}):\n{result.output}"
+        assert "Dynamic axes:" in result.output
+        assert output_file.exists()
+        data = json.loads(output_file.read_text())
+
+        # The benchmarked model exposes a dynamic (None) batch axis; the
+        # remaining channel/spatial dims stay static.
+        input_shapes = data["model_info"]["input_shapes"]
+        assert input_shapes[0][0] is None, f"expected dynamic batch, got {input_shapes}"
+        assert input_shapes[0][1:] == [3, 224, 224]
+        assert data["latency_ms"]["mean"] > 0
+
+        # The ONNX graph ORT actually loaded carries the symbolic batch dim.
+        running = onnx.load(data["benchmark_info"]["running_model_path"])
+        pixel_values = next(i for i in running.graph.input if i.name == "pixel_values")
+        assert pixel_values.type.tensor_type.shape.dim[0].dim_param == "batch"
+
+
 # ===========================================================================
 # GenAI runtime (winml-genai): --device / --ep override
 # ===========================================================================