Skip to content

feat: Add dynamic axes export support#1074

Merged
xieofxie merged 4 commits into
mainfrom
hualxie/dynamic_export
Jul 10, 2026
Merged

feat: Add dynamic axes export support#1074
xieofxie merged 4 commits into
mainfrom
hualxie/dynamic_export

Conversation

@xieofxie

@xieofxie xieofxie commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Add a first-class --dynamic-axes option for winml export.
  • Allow symbolic string dimensions in --input-specs and infer matching dynamic ONNX input axes.
  • Normalize/validate dynamic axes mappings, including JSON string axis keys.
  • Keep static shapes as the QNN-safe default while making requested dynamic batch informational instead of warning-shaped.
  • Document dynamic axes, symbolic input specs, and QNN compatibility caveats.

Fixes #929

Validation

  • uv run ruff check --fix src\winml\modelkit\onnx\io.py src\winml\modelkit\export\config.py src\winml\modelkit\commands\export.py src\winml\modelkit\export\htp\exporter.py tests\unit\export\test_config_validation.py tests\unit\export\test_pytorch_export.py tests\unit\commands\test_export.py
  • uv run pytest tests\unit\export\test_config_validation.py tests\unit\export\test_pytorch_export.py tests\unit\commands\test_export.py

@xieofxie
xieofxie requested a review from a team as a code owner July 8, 2026 02:46
@xieofxie xieofxie changed the title Add dynamic axes export support feat: Add dynamic axes export support Jul 8, 2026

@DingmaomaoBJTU DingmaomaoBJTU left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice, well-scoped addition — the docs, config validation, and the end-to-end symbolic-shape export test all look solid, and the layered --dynamic-axes > --export-config precedence is clear.

A few consistency nits below around the dynamic-batch messaging and symbolic-dim validation. Nothing blocking.

Comment thread src/winml/modelkit/export/config.py
Comment thread src/winml/modelkit/export/config.py Outdated
Comment thread src/winml/modelkit/config/build.py Outdated

@DingmaomaoBJTU DingmaomaoBJTU left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Took a pass over the dynamic axes support. The normalization / merge / inference split in export/config.py is clean, symbolic dims resolving to size 1 for the dummy tensor is a sensible default, and the test coverage is solid. A few non-blocking notes inline: an error-wrapping inconsistency in the --export-config path, a cross-module private-method call, and a small question about the conflict-on-merge UX. Nothing blocking from my side.

Comment thread src/winml/modelkit/commands/export.py Outdated
Comment thread src/winml/modelkit/config/build.py Outdated
Comment thread src/winml/modelkit/export/config.py

@DingmaomaoBJTU DingmaomaoBJTU left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did a full holistic pass over the complete change set (not just the incremental delta) and everything checks out. Approving. ✅

What I verified this round:

  • All three of my previous comments were addressed cleanly — the --export-config error-handling is now consistent with --shape-config/--dynamic-axes, concrete_shape() was promoted to a public method used by both call sites, and the conflict-on-merge behavior is intentional with a clear CLI error message.
  • Symbolic (str) dims don't leak into numeric code paths: the batch_size comparison in __post_init__ guards with isinstance(spec.shape[0], int), dummy generation goes through concrete_shape(), and the metadata embedding (to_dict() + json.dumps) round-trips symbolic shapes correctly.
  • dynamic_axes int-key ↔ JSON string-key normalization round-trips through to_dict/from_dict.
  • _verify_onnx_export and __post_init__ are aligned: requested dynamic batch logs INFO, unexpected dynamic batch still WARNs.
  • Explicit --dynamic-axes supports output tensors too (covered by tests).

Checks: ruff clean on all changed files; full CI green (lint, CodeQL, test commands/models/optim/remaining).

Nice, well-tested feature. Thanks for iterating on the feedback!

@xieofxie
xieofxie merged commit 670ad35 into main Jul 10, 2026
9 checks passed
@xieofxie
xieofxie deleted the hualxie/dynamic_export branch July 10, 2026 01:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: support dynamic axes / dynamic dimensions in winml export

2 participants