@@ -3,13 +3,35 @@ from __future__ import annotations
33from openvino ._pyopenvino import PartialShape
44from openvino ._pyopenvino import Tensor
55from openvino .tools .ovc .cli_parser import InputCutInfo as _InputCutInfo
6+ from openvino .tools .ovc .cli_parser import input_to_input_cut_info
67from openvino .tools .ovc .cli_parser import single_input_to_input_cut_info
78from openvino .tools .ovc .error import Error
89import logging as log
910import numpy as np
1011import pathlib as pathlib
1112import sys as sys
12- __all__ : list [str ] = ['Error' , 'PartialShape' , 'Tensor' , 'extract_input_info_from_example' , 'extract_module_extensions' , 'flatten_inputs' , 'get_pytorch_decoder' , 'get_pytorch_decoder_for_model_on_disk' , 'get_value_from_list_or_dict' , 'log' , 'np' , 'pathlib' , 'prepare_torch_inputs' , 'single_input_to_input_cut_info' , 'sys' , 'to_torch_tensor' , 'update_list_or_dict' ]
13+ __all__ : list [str ] = ['Error' , 'PartialShape' , 'Tensor' , 'extract_input_info_from_example' , 'extract_module_extensions' , 'flatten_inputs' , 'get_pytorch_decoder' , 'get_pytorch_decoder_for_model_on_disk' , 'get_value_from_list_or_dict' , 'input_to_input_cut_info' , 'log' , 'np' , 'pathlib' , 'prepare_torch_inputs' , 'single_input_to_input_cut_info' , 'sys' , 'to_torch_tensor' , 'update_list_or_dict' ]
14+ def _build_dynamic_shapes (inputs , input_specs = None ):
15+ """
16+ Build dynamic_shapes for torch.export.export.
17+
18+ If input_specs (list of _InputCutInfo from the 'input' parameter) is provided
19+ and contains shapes, dimensions marked as -1 (fully dynamic) get Dim.AUTO,
20+ dimensions with min/max constraints (e.g. Dimension(1, 10)) get
21+ Dim("dI_D", min=..., max=...), and fixed dimensions stay static.
22+ When no specs are given returns None so that torch.export.export produces
23+ a fully static graph.
24+
25+ The input_specs list is flat (one spec per leaf tensor), while inputs may
26+ contain nested tuples/lists (e.g. past_key_values). pytree is used to
27+ flatten inputs, pair each leaf with its spec, and then unflatten the
28+ result back into the original structure that torch.export expects.
29+
30+ """
31+ def _export_torch_model (model , inputs , input_specs = None ):
32+ """
33+ Export a torch.nn.Module using torch.export.export with Dim.AUTO dynamic shapes.
34+ """
1335def extract_input_info_from_example (args , inputs ):
1436 ...
1537def extract_module_extensions (args ):
0 commit comments