feat(config): add dynamic export controls to config command

docs/commands/config.md

@@ -23,6 +23,9 @@ $ winml config [options]
 | `--module` | | `TEXT` | *(none)* | Generate configs for every submodule whose class name matches the given string (e.g., `ResNetConvLayer`). The output is a JSON array instead of a single object. |
 | `--config` | `-c` | `PATH` | *(none)* | JSON override file in `WinMLBuildConfig` format. Fields present in this file take precedence over auto-detected values. |
 | `--shape-config` | | `PATH` | *(none)* | JSON file with input shape overrides for dummy input generation. Valid keys by modality — text: `sequence_length`; vision: `height`, `width`, `num_channels`; audio: `feature_size`, `nb_max_frames`, `audio_sequence_length`. |
+| `--input-specs` | | `PATH` | *(none)* | JSON file with input specifications for the HuggingFace export. Fields are patched onto the auto-resolved input tensors *by name* (unlisted inputs and their `dtype`/`value_range` are preserved); unknown names are appended. Only valid for a single HuggingFace export config (not `--module` or composite models). |
+| `--export-config` | | `PATH` | *(none)* | ONNX export configuration JSON (`opset_version`, `do_constant_folding`, etc.) merged into the generated `export` section. Only valid for a single HuggingFace export config (not `--module` or composite models). |
+| `--dynamic-axes` | | `PATH` | *(none)* | JSON dynamic axes mapping for the ONNX export (e.g., `{"input_ids": {"0": "batch", "1": "sequence"}}`). Symbolic string dimensions in `--input-specs` shapes also infer dynamic axes. Only valid for a single HuggingFace export config (not `--module` or composite models). |
 | `--device` | `-d` | `auto\|npu\|gpu\|cpu` | `auto` | Target device. Affects the generated quantization and compilation sub-configs. `auto` leaves those sections unchanged from the kit defaults. |
 | `--ep` | | `TEXT` | *(none)* | Force a specific execution provider (`qnn`, `dml`, `migraphx`, `tensorrt`, `vitisai`, `openvino`, `cpu`). Overrides the device-to-provider mapping. When used without `--device`, the device is inferred from the EP. |
 | `--precision` | `-p` | `TEXT` | `auto` | Target precision: `auto`, `fp32`, `fp16`, `int8`, `int16`, or a mixed format such as `w8a16`. `auto` selects the precision based on the chosen device. |
@@ -81,13 +84,22 @@ Generate a config from an already-exported ONNX file, skipping quantization (com
 $ winml config -m facebook/convnext-tiny-224.onnx --no-quant -o convnext_optim_only.json
 ```
 
+Generate a config with dynamic axes so the exported model accepts a variable batch dimension:
+
+```bash
+$ winml config -m microsoft/resnet-50 --dynamic-axes dynamic_axes.json -o resnet_dyn.json
+```
+
+The `dynamic_axes.json` maps each input to its dynamic dimensions, e.g. `{"pixel_values": {"0": "batch"}}`. Symbolic string dimensions in `--input-specs` shapes (e.g. `{"input_ids": {"shape": ["batch", "sequence"]}}`) infer dynamic axes automatically without a separate `--dynamic-axes` file.
+
 ## Common pitfalls
 
 - **At least one of `-m`, `--model-type`, or `--model-class` is required** — calling `winml config` with none of these three flags raises a usage error immediately.
 - **`auto` precision does not always map to a lower-bit type** — when `--device` is also `auto`, precision stays at the kit default (usually `fp32`). Explicitly pass `--device npu` or `--device gpu` for `auto` precision to resolve to `int8` or `fp16`.
 - **`--module` changes the output shape** — with `--module` the JSON output is an array of configs, not a single object. Scripts that expect a single object will fail to parse this output.
 - **`--trust-remote-code` has security implications** — only use this flag with model repositories you own or explicitly trust; it allows arbitrary Python execution from the remote model card.
 - **Shape overrides in `--shape-config` are modality-specific** — passing a `sequence_length` key for a vision model has no effect. Check the `--help` description for valid keys per modality.
+- **Export controls require a single HuggingFace export graph** — `--input-specs`, `--export-config`, and `--dynamic-axes` are rejected for pre-exported `.onnx` inputs (which set `export` to `null`); use them only when config generates the export section. They are also rejected for composite (multi-component) models and for `--module` (which fan out one config per sub-component / submodule, each with their own export inputs) — generate those per-config outputs first, then edit their export sections individually.
 
 ## See also
 

src/winml/modelkit/commands/config.py

@@ -59,6 +59,27 @@ def _apply_stage_overrides(cfg: Any, *, no_quant: bool, no_compile: bool) -> Non
         cfg.compile = None
 
 
+def _merge_export_overrides(cfg: Any, export_overrides: dict[str, Any]) -> Any:
+    """Apply --export-config/--dynamic-axes/--input-specs onto a generated config.
+
+    Returns ``cfg`` unchanged when no export overrides were supplied. Mirrors the
+    ``build`` command: ``--input-specs`` patches the auto-resolved input tensors
+    by name (preserving unlisted inputs and their dtype/value_range) and any
+    symbolic dims re-derive dynamic axes via ``WinMLExportConfig.__post_init__``.
+    """
+    if not export_overrides:
+        return cfg
+    if cfg.export is None:
+        raise click.UsageError(
+            "--input-specs, --export-config, and --dynamic-axes require a "
+            "HuggingFace export config; they are not supported when the "
+            "generated build config has export=null."
+        )
+    from ..config import merge_export_overrides
+
+    return merge_export_overrides(cfg, export_overrides)
+
+
 @click.command("config")
 @cli_utils.model_option(required=False, optional_message="Optional when --model-type is provided.")
 @click.option(
@@ -99,6 +120,14 @@ def _apply_stage_overrides(cfg: Any, *, no_quant: bool, no_compile: bool) -> Non
     "vision: height, width, num_channels; "
     "audio: feature_size, nb_max_frames, audio_sequence_length.",
 )
+@cli_utils.input_specs_option()
+@cli_utils.export_config_option()
+@cli_utils.dynamic_axes_option(
+    help_text=(
+        "JSON dynamic axes mapping for HuggingFace ONNX export "
+        '(e.g., {"input_ids": {"0": "batch", "1": "sequence"}}).'
+    )
+)
 @cli_utils.device_option(
     required=False,
     optional_message="Affects quant/compile config.",
@@ -140,6 +169,9 @@ def config(
     module: str | None,
     config_file: str | None,
     shape_config_file: str | None,
+    input_specs: Path | None,
+    export_config: Path | None,
+    dynamic_axes: Path | None,
     device: str,
     ep: EPNameOrAlias | None,
     precision: str,
@@ -190,6 +222,10 @@ def config(
         # Vision model with shape overrides ({"height": 224, "width": 224})
         winml config --model-type resnet -t image-classification --shape-config shapes.json
 
+        # Dynamic export controls (mirrors ``winml build`` / ``winml export``)
+        winml config -m bert-base-uncased --dynamic-axes dynamic_axes.json
+        winml config -m bert-base-uncased --input-specs inputs.json --export-config export.json
+
         # Save to file
         winml config -m bert-base-uncased -o config.json
 
@@ -276,6 +312,27 @@ def config(
                 "--module is not supported with ONNX file input. "
                 "Module discovery requires a HuggingFace model."
             )
+        # Export controls (--input-specs/--export-config/--dynamic-axes) target a
+        # single export graph, so they only apply to the plain single-config
+        # HuggingFace path. Reject the multi-graph / no-export paths up front — on
+        # raw flag presence, before loading/validating the JSON — so the error
+        # names the real problem (ONNX input, --module fan-out, composite) instead
+        # of a downstream "Invalid export configuration", and we skip needless I/O.
+        _export_flags_given = bool(input_specs or export_config or dynamic_axes)
+        if hf_model and _hf_is_onnx and _export_flags_given:
+            raise click.UsageError(
+                "--input-specs, --export-config, and --dynamic-axes are only "
+                "supported when generating a HuggingFace export config, not "
+                "pre-exported ONNX files."
+            )
+        if module and _export_flags_given:
+            raise click.UsageError(
+                "--input-specs, --export-config, and --dynamic-axes are not "
+                "supported with --module, which generates one config per matched "
+                "submodule. Generate the per-module configs first, then edit their "
+                "export sections individually."
+            )
+
         config_obj: WinMLBuildConfig | None = None
         output_data: dict[str, Any] | list[Any]
         if hf_model and _hf_is_onnx:
@@ -306,6 +363,20 @@ def config(
                 hf_model, model_type, task, trust_remote_code=trust_remote_code
             )
             if pipeline_components:
+                # Export controls target a single export graph; a composite model
+                # has one export per sub-component with distinct inputs. Reject on
+                # raw flag presence — before loading/validating the JSON — so the
+                # composite-specific error wins (mirroring the ONNX path) instead
+                # of a downstream "Invalid export configuration". config never fans
+                # these overrides out across heterogeneous components.
+                if _export_flags_given:
+                    raise click.UsageError(
+                        "--input-specs, --export-config, and --dynamic-axes are not "
+                        "supported for composite (multi-component) models, whose "
+                        "sub-components each have their own export inputs. Generate "
+                        "the per-component configs first, then edit their export "
+                        "sections individually."
+                    )
                 # composite model: generate one config per sub-component
                 _generate_pipeline_configs(
                     pipeline_components,
@@ -327,6 +398,16 @@ def config(
                 )
                 return
 
+            # Load export CLI overrides now that the multi-graph paths (ONNX,
+            # --module, composite) have all been rejected — the single HF config
+            # path below is the only one that applies them. Returned sparse so
+            # unspecified fields don't clobber auto-detected values.
+            export_overrides = cli_utils.load_export_overrides(
+                export_config=export_config,
+                input_specs=input_specs,
+                dynamic_axes=dynamic_axes,
+            )
+
             # Generate config(s). The ``module: str | None`` overload of
             # generate_hf_build_config returns WinMLBuildConfig | list[...],
             # which isinstance(result, list) narrows for the branches below.
@@ -345,15 +426,18 @@ def config(
                 ep=ep,
             )
             if isinstance(result, list):
-                configs = result
+                # --module + export overrides is rejected up front, so
+                # export_overrides is empty here; emit the submodule configs as
+                # generated without fanning any overrides across them.
+                configs = list(result)
                 for cfg in configs:
                     _apply_stage_overrides(cfg, no_quant=not quant, no_compile=no_compile)
                 output_data = [cfg.to_dict() for cfg in configs]
                 _n_modules = len(configs)
                 # Use first config for display metadata
                 config_obj = configs[0] if configs else None
             else:
-                config_obj = result
+                config_obj = _merge_export_overrides(result, export_overrides)
                 configs = []
                 _apply_stage_overrides(config_obj, no_quant=not quant, no_compile=no_compile)
                 output_data = config_obj.to_dict()
@@ -407,6 +491,21 @@ def config(
                 f"   \U0001f4c1 [bold]Shape config:[/bold] "
                 f"{_shape_config_file}  [green]\u2713[/green]"
             )
+        if input_specs:
+            console.print(
+                f"   \U0001f4c1 [bold]Input specs:[/bold]  "
+                f"{input_specs.name}  [green]\u2713[/green]"
+            )
+        if export_config:
+            console.print(
+                f"   \U0001f4c1 [bold]Export config:[/bold] "
+                f"{export_config.name}  [green]\u2713[/green]"
+            )
+        if dynamic_axes:
+            console.print(
+                f"   \U0001f4c1 [bold]Dynamic axes:[/bold] "
+                f"{dynamic_axes.name}  [green]\u2713[/green]"
+            )
 
         console.print()
 

tests/unit/commands/test_build.py

@@ -88,6 +88,25 @@ def mock_task_model_compatibility_validator():
         yield
 
 
+@pytest.fixture(autouse=True)
+def mock_composite_resolution():
+    """Default composite detection to "not composite" without any network call.
+
+    ``build`` calls ``resolve_composite_components`` before dispatching, which
+    resolves the model config from HuggingFace Hub (``AutoConfig.from_pretrained``)
+    over the network. Left live in these CLI plumbing tests, repeated metadata
+    requests get throttled and ``huggingface_hub``'s retry backoff sleeps — making
+    a later, unrelated test appear to hang. Returning ``None`` keeps the plain
+    single-build path. Composite-specific tests override this with their own
+    ``patch(...)``, which nests inside and wins for the duration of that test.
+    """
+    with patch(
+        "winml.modelkit.loader.resolution.resolve_composite_components",
+        return_value=None,
+    ):
+        yield
+
+
 @pytest.fixture
 def runner() -> CliRunner:
     """Create a CLI test runner."""