From 44a729241e105b5abe502a1e4f7005ab2d759e2b Mon Sep 17 00:00:00 2001 From: danielliuce Date: Mon, 12 Jan 2026 10:46:45 -0500 Subject: [PATCH 1/6] support input shapes export --- optimum/commands/export/openvino.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/optimum/commands/export/openvino.py b/optimum/commands/export/openvino.py index cd3280189e..09c9548198 100644 --- a/optimum/commands/export/openvino.py +++ b/optimum/commands/export/openvino.py @@ -294,6 +294,16 @@ def parse_args_openvino(parser: "ArgumentParser"): type=json.loads, help=("Any kwargs passed to the model forward, or used to customize the export for a given model."), ) + optional_group.add_argument( + "--input-shapes", + type=json.loads, + default=None, + help=( + "Override the default shapes used during export. Provide shapes as JSON, e.g., " + '\'{{"batch_size": 1, "sequence_length": 128}}\' for static shapes. ' + "Common shape parameters include: batch_size, sequence_length, height, width, num_channels." + ), + ) def no_compression_parameter_provided(args): @@ -464,7 +474,6 @@ def run(self): output = Path(self.args.output) try: - # TODO : add input shapes main_export( model_name_or_path=self.args.model, output=output, @@ -479,7 +488,7 @@ def run(self): library_name=library_name, variant=self.args.variant, model_kwargs=self.args.model_kwargs, - # **input_shapes, + **(self.args.input_shapes or {}), ) if apply_main_quantize: _main_quantize( From 1bc03f1baffa0e4ceba095b08e1258e01826f5e0 Mon Sep 17 00:00:00 2001 From: danielliuce Date: Mon, 12 Jan 2026 12:08:01 -0500 Subject: [PATCH 2/6] propagate static input shapes downstream --- optimum/exporters/openvino/convert.py | 6 ++++-- optimum/exporters/openvino/utils.py | 8 ++++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index 0ed08a5f2f..ca94efe0c7 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -361,8 +361,10 @@ def export_pytorch( logger.info(f"\t- {override_config_key} -> {override_config_value}") setattr(model.config, override_config_key, override_config_value) + static_shapes = True if input_shapes is None: input_shapes = {} # will use the defaults from DEFAULT_DUMMY_SHAPES + static_shapes = False # User did not explicit input shapes, produce dynamic shapes in the model # Check that inputs match, and order them properly dummy_inputs = config.generate_dummy_inputs(framework="pt", **input_shapes) @@ -411,7 +413,7 @@ def ts_patched_forward(*args, **kwargs): with patcher: check_dummy_inputs_are_allowed(model, dummy_inputs) - input_info = _get_input_info(model, config, dummy_inputs) + input_info = _get_input_info(model, config, dummy_inputs, static_shapes=static_shapes) torch_export = os.getenv("OPENVINO_DYNAMO_EXPORT", "false").lower() == "true" if torch_export: if hasattr(torch.ops, "_prepare_4d_causal_attention_mask_for_sdpa"): @@ -633,7 +635,7 @@ def export_from_model( # Get the shapes to be used to generate dummy inputs input_shapes = {} for input_name in DEFAULT_DUMMY_SHAPES.keys(): - if input_name in ["height", "width"]: + if input_name in ["height", "width"] and kwargs_shapes is None: # use H and W from generator defaults continue input_shapes[input_name] = ( diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index eea8af6b6d..44726a6e3e 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -87,7 +87,10 @@ def flattenize_inputs(inputs: List[Any]): def _get_input_info( - model: Union["PreTrainedModel", "ModelMixin"], config: OnnxConfig, dummy_inputs: Dict[str, Any] + model: Union["PreTrainedModel", "ModelMixin"], + config: OnnxConfig, + dummy_inputs: Dict[str, Any], + static_shapes: bool = False ) -> List[InputInfo]: sig = inspect.signature(model.forward) if hasattr(model, "forward") else inspect.signature(model.call) inputs = config.ordered_inputs(model) @@ -105,7 +108,8 @@ def _get_input_info( example = flatten_inputs[i] type = get_element_type(example.cpu().numpy().dtype) shape = PartialShape(example.shape) - if name in inputs: + # Set dynamic axes from config, unless user requested static shapes + if name in inputs and not static_shapes: named_dims = inputs[name] for idx, dim_name in named_dims.items(): if dim_name in name_to_symbol: From 1def1bc80ae4a4049efd92dfd9c1563d2b88c1af Mon Sep 17 00:00:00 2001 From: danielliuce Date: Tue, 13 Jan 2026 12:25:19 -0500 Subject: [PATCH 3/6] fix formatting --- optimum/exporters/openvino/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index 44726a6e3e..b11d9ee9a4 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -90,7 +90,7 @@ def _get_input_info( model: Union["PreTrainedModel", "ModelMixin"], config: OnnxConfig, dummy_inputs: Dict[str, Any], - static_shapes: bool = False + static_shapes: bool = False, ) -> List[InputInfo]: sig = inspect.signature(model.forward) if hasattr(model, "forward") else inspect.signature(model.call) inputs = config.ordered_inputs(model) From 825640128b6041143285b7cfb48d9717d18306b3 Mon Sep 17 00:00:00 2001 From: Daniel Liu Date: Wed, 14 Jan 2026 17:16:42 -0500 Subject: [PATCH 4/6] implement more intuitive command line arguments for static shapes --- optimum/commands/export/openvino.py | 84 +++++++++++++++++++++++++-- optimum/exporters/openvino/convert.py | 5 +- 2 files changed, 81 insertions(+), 8 deletions(-) diff --git a/optimum/commands/export/openvino.py b/optimum/commands/export/openvino.py index 09c9548198..0edd678fe1 100644 --- a/optimum/commands/export/openvino.py +++ b/optimum/commands/export/openvino.py @@ -22,6 +22,7 @@ from optimum.commands.base import BaseOptimumCLICommand, CommandInfo from optimum.utils.constant import ALL_TASKS +from optimum.utils import DEFAULT_DUMMY_SHAPES logger = logging.getLogger(__name__) @@ -294,16 +295,82 @@ def parse_args_openvino(parser: "ArgumentParser"): type=json.loads, help=("Any kwargs passed to the model forward, or used to customize the export for a given model."), ) + doc_input = "to use in the example input given to the OpenVINO export." optional_group.add_argument( - "--input-shapes", - type=json.loads, + "--batch_size", + type=int, + default=None, + help=f"Batch size {doc_input}", + ) + optional_group.add_argument( + "--sequence_length", + type=int, + default=None, + help=f"Text tasks only. Sequence length {doc_input}", + ) + optional_group.add_argument( + "--num_choices", + type=int, + default=None, + help=f"Text tasks only. Num choices {doc_input}", + ) + optional_group.add_argument( + "--width", + type=int, + default=None, + help=f"Image tasks only. Width {doc_input}", + ) + optional_group.add_argument( + "--height", + type=int, + default=None, + help=f"Image tasks only. Height {doc_input}", + ) + optional_group.add_argument( + "--num_channels", + type=int, + default=None, + help=f"Image tasks only. Number of channels {doc_input}", + ) + optional_group.add_argument( + "--feature_size", + type=int, + default=None, + help=f"Audio tasks only. Feature size {doc_input}", + ) + optional_group.add_argument( + "--nb_max_frames", + type=int, + default=None, + help=f"Audio tasks only. Maximum number of frames {doc_input}", + ) + optional_group.add_argument( + "--audio_sequence_length", + type=int, + default=None, + help=f"Audio tasks only. Audio sequence length {doc_input}", + ) + optional_group.add_argument( + "--point_batch_size", + type=int, default=None, help=( - "Override the default shapes used during export. Provide shapes as JSON, e.g., " - '\'{{"batch_size": 1, "sequence_length": 128}}\' for static shapes. ' - "Common shape parameters include: batch_size, sequence_length, height, width, num_channels." + "For Segment Anything. It corresponds to how many segmentation masks we want the model to predict per " + "input point." ), ) + optional_group.add_argument( + "--nb_points_per_image", + type=int, + default=None, + help="For Segment Anything. It corresponds to the number of points per segmentation masks.", + ) + optional_group.add_argument( + "--visual_seq_length", + type=int, + default=None, + help="Visual sequence length", + ) def no_compression_parameter_provided(args): @@ -474,6 +541,11 @@ def run(self): output = Path(self.args.output) try: + # Get the shapes to be used to generate dummy inputs + input_shapes = {} + for input_name in DEFAULT_DUMMY_SHAPES: + if hasattr(self.args, input_name) and getattr(self.args, input_name) is not None: + input_shapes[input_name] = getattr(self.args, input_name) main_export( model_name_or_path=self.args.model, output=output, @@ -488,7 +560,7 @@ def run(self): library_name=library_name, variant=self.args.variant, model_kwargs=self.args.model_kwargs, - **(self.args.input_shapes or {}), + **input_shapes ) if apply_main_quantize: _main_quantize( diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index ca94efe0c7..411f797a54 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -362,8 +362,8 @@ def export_pytorch( setattr(model.config, override_config_key, override_config_value) static_shapes = True - if input_shapes is None: - input_shapes = {} # will use the defaults from DEFAULT_DUMMY_SHAPES + if input_shapes["static_shapes"] == False: + input_shapes = {} # will dynamic shapes in the OpenVINO IR static_shapes = False # User did not explicit input shapes, produce dynamic shapes in the model # Check that inputs match, and order them properly @@ -641,6 +641,7 @@ def export_from_model( input_shapes[input_name] = ( kwargs_shapes[input_name] if input_name in kwargs_shapes else DEFAULT_DUMMY_SHAPES[input_name] ) + input_shapes["static_shapes"] = True if kwargs_shapes else False if library_name == "open_clip": custom_architecture = True From 7b00e498c8b66c74e10038f6785f5051f33d0997 Mon Sep 17 00:00:00 2001 From: Daniel Liu Date: Wed, 14 Jan 2026 17:28:27 -0500 Subject: [PATCH 5/6] fix comment --- optimum/exporters/openvino/convert.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index 411f797a54..88ca40fd1c 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -363,7 +363,7 @@ def export_pytorch( static_shapes = True if input_shapes["static_shapes"] == False: - input_shapes = {} # will dynamic shapes in the OpenVINO IR + input_shapes = {} # will use dynamic shapes in the OpenVINO IR static_shapes = False # User did not explicit input shapes, produce dynamic shapes in the model # Check that inputs match, and order them properly From 490e69062eed230ffac42fc3648d3830df99c0a8 Mon Sep 17 00:00:00 2001 From: Daniel Liu Date: Wed, 14 Jan 2026 17:29:54 -0500 Subject: [PATCH 6/6] fix formatting --- optimum/commands/export/openvino.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/commands/export/openvino.py b/optimum/commands/export/openvino.py index 0edd678fe1..012e813283 100644 --- a/optimum/commands/export/openvino.py +++ b/optimum/commands/export/openvino.py @@ -560,7 +560,7 @@ def run(self): library_name=library_name, variant=self.args.variant, model_kwargs=self.args.model_kwargs, - **input_shapes + **input_shapes, ) if apply_main_quantize: _main_quantize(