feat: support dynamic axes / dynamic dimensions in `winml export`

[xieofxie]

Summary

The winml export command (src/winml/modelkit/commands/export.py) has no first-class support for dynamic axes / dynamic dimensions in the exported ONNX. Everything the CLI surfaces is aimed at producing static shapes. Dynamic axes are only reachable by smuggling a raw dynamic_axes dict through an --export-config JSON, and the toolchain actively warns against it. We should decide on, and implement, proper dynamic-dimension support (e.g. dynamic batch / sequence length).

Current state (findings)

The CLI exposes only static-shape controls

• No --dynamic-axes (or equivalent) flag exists on the export command.
• --shape-config (export.py:126) takes concrete integer overrides like {"sequence_length": 2048, "height": 640} — it pins symbolic dims to fixed values (the opposite of dynamic).
• --input-specs shapes are parsed as a fixed tuple(spec["shape"]) (export.py:334) — concrete integers only, no symbolic axis names.

The capability exists underneath but is not wired to the CLI

WinMLExportConfig already has a dynamic_axes field that is forwarded to torch.onnx.export:

• src/winml/modelkit/export/config.py:83 — dynamic_axes: dict[str, dict[int, str]] | None = None
• src/winml/modelkit/export/config.py:301 — from_dict reads data.get("dynamic_axes")
• src/winml/modelkit/export/htp/exporter.py:442 — if export_config.dynamic_axes: onnx_kwargs["dynamic_axes"] = export_config.dynamic_axes

Because export.py loads the --export-config JSON into config_kwargs and builds the config via WinMLExportConfig.from_dict(config_kwargs) (export.py:353, export.py:369), a user can currently enable dynamic dims only by hand-writing a "dynamic_axes": {...} key in an --export-config JSON file. There is no dedicated, validated, discoverable option.

The HTP/QNN path deliberately discourages dynamic dims

The default export path is built around static dimensions on purpose:

• exporter.py:76-77 — "dynamic_axes: Not set (defaults to None = static dimensions). This prevents dynamic batch which causes MatMulAddFusion failure."
• exporter.py:369-384 — after export it inspects the graph and warns if it finds a symbolic (dim_param) batch dimension, recommending removal of dynamic_axes.
• config.py:144-148 — logs a warning if dynamic_axes marks axis 0 (batch) dynamic.

Why this matters

Some downstream scenarios (variable batch size, variable sequence length, variable image resolution) require dynamic dimensions in the exported ONNX. Today users must reverse-engineer the dynamic_axes schema and bypass the intended CLI surface, with no validation and conflicting warnings.

Proposed scope (to refine)

• Decide the supported surface: a dedicated --dynamic-axes option and/or a dynamic field per input in --input-specs / --shape-config.
• Allow symbolic dimension names (strings) in InputTensorSpec.shape instead of integers-only.
• Reconcile with the HTP/QNN static-shape requirement: when does dynamic break fusions (e.g. MatMulAddFusion), and should it be blocked / gated per EP?
• Update the warnings in exporter.py / config.py so intentional dynamic export isn't flagged as an error.
• Add tests covering a dynamic-axes export (none currently exercise this path through the CLI).
• Document the feature and its compatibility caveats.

Open questions

• Should dynamic axes be opt-in per-EP (e.g. allowed for CPU/CUDA, blocked/warned for QNN)?
• What is the desired UX — explicit dynamic_axes dict, or higher-level flags like --dynamic-batch / --dynamic-seq-len?