feat: Add dynamic axes export support#1074
Conversation
DingmaomaoBJTU
left a comment
There was a problem hiding this comment.
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.
DingmaomaoBJTU
left a comment
There was a problem hiding this comment.
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.
DingmaomaoBJTU
left a comment
There was a problem hiding this comment.
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-configerror-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: thebatch_sizecomparison in__post_init__guards withisinstance(spec.shape[0], int), dummy generation goes throughconcrete_shape(), and the metadata embedding (to_dict()+json.dumps) round-trips symbolic shapes correctly. dynamic_axesint-key ↔ JSON string-key normalization round-trips throughto_dict/from_dict._verify_onnx_exportand__post_init__are aligned: requested dynamic batch logs INFO, unexpected dynamic batch still WARNs.- Explicit
--dynamic-axessupports 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!
Summary
--dynamic-axesoption forwinml export.--input-specsand infer matching dynamic ONNX input axes.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.pyuv run pytest tests\unit\export\test_config_validation.py tests\unit\export\test_pytorch_export.py tests\unit\commands\test_export.py